wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-09-11 12:43:20 +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/batch_norm_model.py
+5 -4
View File
@@ -174,17 +174,18 @@ class TorchBatchNormModel(TorchModelV2, nn.Module):
activation_fn=None)
self._hidden_layers = nn.Sequential(*layers)
self._hidden_out = None
@override(ModelV2)
def forward(self, input_dict, state, seq_lens):
# Set the correct train-mode for our hidden module (only important
# b/c we have some batch-norm layers).
self._hidden_layers.train(mode=input_dict["is_training"])
hidden_out = self._hidden_layers(input_dict["obs"])
logits = self._logits(hidden_out)
self._value_out = self._value_branch(hidden_out)
self._hidden_out = self._hidden_layers(input_dict["obs"])
logits = self._logits(self._hidden_out)
return logits, []
@override(ModelV2)
def value_function(self):
return torch.reshape(self._value_out, [-1])
assert self._hidden_out is not None, "must call forward first!"
return torch.reshape(self._value_branch(self._hidden_out), [-1])