wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-09-09 11:32:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] Make PyTorch Model forward pass faster in vf-case. (#8422)

Browse Source
This commit is contained in:
Sven Mika
2020-05-14 10:15:50 +02:00
committed by GitHub
parent 212f78f735
commit 5f4c196fed
11 changed files with 55 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/examples/models/shared_weights_model.py
+5 -4
View File
@@ -110,15 +110,16 @@ class TorchSharedWeightsModel(TorchModelV2, nn.Module):
# Non-shared final layer.
self.last_layer = SlimFC(32, self.num_outputs, activation_fn=nn.ReLU)
self.vf = SlimFC(32, 1, activation_fn=None)
self._output = None
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
out = self.first_layer(input_dict["obs"])
out = TORCH_GLOBAL_SHARED_LAYER(out)
model_out = self.last_layer(out)
self._value_out = self.vf(out)
self._output = TORCH_GLOBAL_SHARED_LAYER(out)
model_out = self.last_layer(self._output)
return model_out, []
@override(ModelV2)
def value_function(self):
return torch.reshape(self._value_out, [-1])
assert self._output is not None, "must call forward first!"
return torch.reshape(self.vf(self._output), [-1])