Specify which layers to restore

Shawan_Mohammed · March 27, 2023, 12:16am

How to restore only a part of the neural network in rllib? Does “custom_model_config” have a layer wise restore parameter which I can use?

mannyv · March 27, 2023, 1:35am

Hi @Shawan_Mohammed,

Welcome to the forum.

This kind of feature is not built into rllib. One option you have is to create your own custom policy and override the set_weights method.

Here is the current set_weights method for torch as an example.

github.com

ray-project/ray/blob/10d12f8220d133b76b33ca6b5a1eeb72636f2b4c/rllib/policy/torch_policy.py#L734


      
              )
              return data
          
          
@override(Policy)
          @DeveloperAPI
          def get_weights(self) -> ModelWeights:
              return {k: v.cpu().detach().numpy() for k, v in self.model.state_dict().items()}
          
          
@override(Policy)
          @DeveloperAPI
          def set_weights(self, weights: ModelWeights) -> None:
              weights = convert_to_torch_tensor(weights, device=self.device)
              self.model.load_state_dict(weights)
          
          
@override(Policy)
          @DeveloperAPI
          def is_recurrent(self) -> bool:
              return self._is_recurrent
          
          
@override(Policy)
          @DeveloperAPI

Topic		Replies	Views
How to restore part of weights from trained model into another model before training? RLlib	0	272	April 12, 2021
How do I copy the model? RLlib	2	434	June 28, 2021
Issue with Custom PyTorch Model in Ray RLlib RLlib	0	282	November 3, 2023
Will the added model be saved and loaded? RLlib	6	285	June 30, 2023
Default model architecture for tf and torch on some algorithms seem to be different RLlib	1	366	March 10, 2023

Specify which layers to restore

Related topics