| 69 | |
| 70 | |
| 71 | class MLP_Block(nn.Module): |
| 72 | def __init__(self, |
| 73 | in_dim: int, |
| 74 | out_dim: int, |
| 75 | net_norm: str = 'none', |
| 76 | activation_function: str = 'Gelu', |
| 77 | dropout: float = 0.0, |
| 78 | residual: bool = False |
| 79 | ): |
| 80 | super().__init__() |
| 81 | layer = [nn.Linear(in_dim, out_dim, bias=net_norm != 'batch_norm')] |
| 82 | if net_norm != 'none': |
| 83 | layer += [get_normalization_layer(net_norm, out_dim)] |
| 84 | layer += [get_activation_function(activation_function, functional=False)] |
| 85 | if dropout > 0: |
| 86 | layer += [nn.Dropout(dropout)] |
| 87 | self.layer = nn.Sequential(*layer) |
| 88 | self.residual = residual |
| 89 | if in_dim != out_dim: |
| 90 | self.residual = False |
| 91 | elif residual: |
| 92 | print('MLP block with residual!') |
| 93 | |
| 94 | def forward(self, X: Tensor) -> Tensor: |
| 95 | X_out = self.layer(X) |
| 96 | if self.residual: |
| 97 | X_out += X |
| 98 | return X_out |