-
Notifications
You must be signed in to change notification settings - Fork 0
/
residual_mlp_blocks.py
213 lines (181 loc) · 7.15 KB
/
residual_mlp_blocks.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
"""
Standard multi-layer perceptron dynamics model.
Author: Ian Char
"""
from typing import Sequence, Tuple, Dict, Any, Optional, Callable
import torch
from torchmetrics import ExplainedVariance
from dynamics_toolbox.models.pl_models.abstract_pl_model import AbstractPlModel
import dynamics_toolbox.constants.activations as activations
import dynamics_toolbox.constants.losses as losses
from dynamics_toolbox.utils.misc import get_architecture
from dynamics_toolbox.utils.pytorch.activations import get_activation
from dynamics_toolbox.utils.pytorch.losses import get_regression_loss
from dynamics_toolbox.utils.pytorch.modules.fc_network import FCNetwork
class ResidualMLPBlocks(AbstractPlModel):
"""Fully connected network for dynamics."""
def __init__(
self,
input_dim: int,
output_dim: int,
embed_dim: int,
num_layers_per_block: int,
num_blocks: int,
learning_rate: float = 3e-4,
hidden_activation: str = activations.RELU,
loss_type: str = losses.MSE,
weight_decay: Optional[float] = 1e-3,
**kwargs,
):
"""Constructor.
Args:
input_dim: The input dimension.
output_dim: The output dimension.
learning_rate: The learning rate for the network.
num_layers: The number of hidden layers in the MLP.
layer_size: The size of each hidden layer in the MLP.
architecture: The architecture of the MLP described as a
a string of underscore separated ints e.g. 256_100_64.
If provided, this overrides num_layers and layer_sizes.
hidden_activation: Activation to use.
loss_type: The name of the loss function to use.
weight_decay: The weight decay for the optimizer.
"""
super().__init__(input_dim, output_dim, **kwargs)
if num_blocks < 3:
raise ValueError('Require number of blocks to be greater than 3.')
self.num_blocks = num_blocks
hidden_sizes = get_architecture(num_layers_per_block, embed_dim, None)
for bnum in range(num_blocks):
indim = embed_dim if bnum > 0 else input_dim
outdim = embed_dim if bnum < num_blocks - 1 else output_dim
setattr(self, f'block_{bnum}', FCNetwork(
input_dim=indim,
output_dim=outdim,
hidden_sizes=hidden_sizes,
hidden_activation=get_activation(hidden_activation),
))
self._learning_rate = learning_rate
self._weight_decay = weight_decay
self._sample_mode = ''
self._loss_function = get_regression_loss(loss_type)
self._loss_type = loss_type
# TODO: In the future we may want to pass this in as an argument.
self._metrics = {
'EV': ExplainedVariance(),
'IndvEV': ExplainedVariance('raw_values'),
}
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward function for network
Args:
x: The input to the network.
Returns:
The output of the network.
"""
x = self.block_0(x)
for bnum in range(1, self.num_blocks - 1):
x = x + getattr(self, f'block_{bnum}')(x)
x = getattr(self, f'block_{self.num_blocks - 1}')(x)
return x
def single_sample_output_from_torch(
self,
net_in: torch.Tensor
) -> Tuple[torch.Tensor, Dict[str, Any]]:
"""Get the output for a single sample in the model.
Args:
net_in: The input for the network.
Returns:
The predictions for a single function sample
"""
with torch.no_grad():
predictions = self.forward(net_in)
info = {'predictions': predictions}
return predictions, info
def multi_sample_output_from_torch(
self,
net_in: torch.Tensor
) -> Tuple[torch.Tensor, Dict[str, Any]]:
"""Get the output where each input is assumed to be from a different sample.
Args:
net_in: The input for the network.
Returns:
The deltas for next states and dictionary of info.
"""
return self.single_sample_output_from_torch(net_in)
def get_net_out(self, batch: Sequence[torch.Tensor]) -> Dict[str, torch.Tensor]:
"""Get the output of the network and organize into dictionary.
Args:
batch: The batch passed to the network.
Returns:
Dictionary of name to tensor.
"""
xi, _ = batch
output = self.forward(xi)
return {'prediction': output}
def loss(
self,
net_out: Dict[str, torch.Tensor],
batch: Sequence[torch.Tensor],
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
"""Compute the loss function.
Args:
net_out: The output of the network.
batch: The batch passed into the network.
Returns:
The loss and a dictionary of other statistics.
"""
_, yi = batch
loss = self._loss_function(net_out['prediction'], yi)
stats = {'loss': loss.item()}
return loss, stats
@property
def sample_mode(self) -> str:
"""The sample mode is the method that in which we get next state."""
return self._sample_mode
@sample_mode.setter
def sample_mode(self, mode: str) -> None:
"""Set the sample mode to the appropriate mode."""
self._sample_mode = mode
@property
def input_dim(self) -> int:
"""The sample mode is the method that in which we get next state."""
return self._hparams.input_dim
@property
def output_dim(self) -> int:
"""The sample mode is the method that in which we get next state."""
return self._hparams.output_dim
@property
def metrics(self) -> Dict[str, Callable[[torch.Tensor], torch.Tensor]]:
"""Get the list of metric functions to compute."""
return self._metrics
@property
def learning_rate(self) -> float:
"""Get the learning rate."""
return self._learning_rate
@property
def weight_decay(self) -> float:
"""Get the weight decay."""
return self._weight_decay
def _get_test_and_validation_metrics(
self,
net_out: Dict[str, torch.Tensor],
batch: Sequence[torch.Tensor],
) -> Dict[str, torch.Tensor]:
"""Compute additional metrics to be used for validation/test only.
Args:
net_out: The output of the network.
batch: The batch passed into the network.
Returns:
A dictionary of additional metrics.
"""
to_return = {}
pred = net_out['prediction']
_, yi = batch
for metric_name, metric in self._metrics.items():
metric_value = metric(pred, yi)
if len(metric_value.shape) > 0:
for dim_idx, metric_v in enumerate(metric_value):
to_return[f'{metric_name}_dim{dim_idx}'] = metric_v
else:
to_return[metric_name] = metric_value
return to_return