| | import torch.nn as nn |
| |
|
| | __all__ = ['SharedMLP'] |
| |
|
| |
|
| | class SharedMLP(nn.Module): |
| | def __init__(self, in_channels, out_channels, dim=1, device='cuda'): |
| | super().__init__() |
| | |
| | if dim == 1: |
| | conv = nn.Conv1d |
| | bn = nn.InstanceNorm1d |
| | elif dim == 2: |
| | conv = nn.Conv2d |
| | bn = nn.InstanceNorm1d |
| | else: |
| | raise ValueError |
| | if not isinstance(out_channels, (list, tuple)): |
| | out_channels = [out_channels] |
| | layers = [] |
| | for oc in out_channels: |
| | layers.extend( |
| | [ |
| | conv(in_channels, oc, 1, device=device), |
| | bn(oc, device=device), |
| | nn.ReLU(True), |
| | ]) |
| | in_channels = oc |
| | self.layers = nn.Sequential(*layers) |
| |
|
| | def forward(self, inputs): |
| | if isinstance(inputs, (list, tuple)): |
| | return (self.layers(inputs[0]), *inputs[1:]) |
| | else: |
| | return self.layers(inputs) |
| |
|
