How to save weights in tensorflow's form?

gavin · August 1, 2021, 11:29am

Hello everyone!

I am working on flow, a computational framework for reinforcement learning in traffic control, which is developed by Berkeley too. In the tutorial, i have already known that the trained neural network models’ parameter can be saved into a ray[rllib] form and rerun by the following code:

trials = run_experiments({
    flow_params["exp_tag"]: {
        "run": alg_run,
        "env": gym_name,
        "config": {
            **config
        },
        "restore": "/ray_results/experiment/dir/checkpoint_50/checkpoint-50"
        "checkpoint_freq": 1,
        "checkpoint_at_end": True,
        "max_failures": 999,
        "stop": {
            "training_iteration": 1,
        },
    },
})

By running this code, i can restart from a checkpoint and continue transfer learning.
However, i am wondering if i can save the weights in a tensorflow’s form so that i can use the model in other simulator (environment).

Thank you!

mannyv · August 1, 2021, 12:14pm

Hi @gavin,

Welcome to the forums. I think calling either export_policy_model on the trainer or export_model model on the policy should do the trick for you. There is an option in those methods to get the model in onnx format. Here is the export_model method for tf.

github.com

ray-project/ray/blob/6703091cdc9f440c886c8c33029bfd5631814c29/rllib/policy/tf_policy.py#L551

    
      
              # Set exploration's state.
              if hasattr(self, "exploration") and "_exploration_state" in state:
                  self.exploration.set_state(
                      state=state["_exploration_state"], sess=self.get_session())
          
          
    # Set the Policy's (NN) weights.
              super().set_state(state)
          
          
@override(Policy)
          @DeveloperAPI
          def export_model(self, export_dir: str,
                           onnx: Optional[int] = None) -> None:
              """Export tensorflow graph to export_dir for serving."""
          
          
    if onnx:
                  try:
                      import tf2onnx
                  except ImportError as e:
                      raise RuntimeError(
                          "Converting a TensorFlow model to ONNX requires "
                          "`tf2onnx` to be installed. Install with "

gavin · August 1, 2021, 2:55pm

Appreciate your answer!

Topic		Replies	Views
Will the added model be saved and loaded? RLlib	6	285	June 30, 2023
How do I copy the model? RLlib	2	434	June 28, 2021
Easiest way with example to save and load models RLlib	2	1153	August 23, 2021
Default model architecture for tf and torch on some algorithms seem to be different RLlib	1	366	March 10, 2023
Some issues happened when I test the method of 'loss' in my 'Policy' class RLlib	0	296	December 7, 2023

How to save weights in tensorflow's form?

Related topics