MLP layers Parameters ---------- input_size: int dim of inputs layer_sizes: list dim of hidden layers output_size: int dim of outputs init_zero: bool zero initialization for the last layer (including w and b). This can provide stable ze
(
input_size,
layer_sizes,
output_size,
output_activation=nn.Identity,
activation=nn.ReLU,
momentum=0.1,
init_zero=False,
)
| 8 | |
| 9 | |
| 10 | def mlp( |
| 11 | input_size, |
| 12 | layer_sizes, |
| 13 | output_size, |
| 14 | output_activation=nn.Identity, |
| 15 | activation=nn.ReLU, |
| 16 | momentum=0.1, |
| 17 | init_zero=False, |
| 18 | ): |
| 19 | """MLP layers |
| 20 | Parameters |
| 21 | ---------- |
| 22 | input_size: int |
| 23 | dim of inputs |
| 24 | layer_sizes: list |
| 25 | dim of hidden layers |
| 26 | output_size: int |
| 27 | dim of outputs |
| 28 | init_zero: bool |
| 29 | zero initialization for the last layer (including w and b). |
| 30 | This can provide stable zero outputs in the beginning. |
| 31 | """ |
| 32 | sizes = [input_size] + layer_sizes + [output_size] |
| 33 | layers = [] |
| 34 | for i in range(len(sizes) - 1): |
| 35 | if i < len(sizes) - 2: |
| 36 | act = activation |
| 37 | layers += [nn.Linear(sizes[i], sizes[i + 1]), |
| 38 | nn.BatchNorm1d(sizes[i + 1], momentum=momentum), |
| 39 | act()] |
| 40 | else: |
| 41 | act = output_activation |
| 42 | layers += [nn.Linear(sizes[i], sizes[i + 1]), |
| 43 | act()] |
| 44 | |
| 45 | if init_zero: |
| 46 | layers[-2].weight.data.fill_(0) |
| 47 | layers[-2].bias.data.fill_(0) |
| 48 | |
| 49 | return nn.Sequential(*layers) |
| 50 | |
| 51 | |
| 52 | def conv3x3(in_channels, out_channels, stride=1): |