File size: 1,057 Bytes
491eded |
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 |
import torch.nn as nn
__all__ = ['SharedMLP']
class SharedMLP(nn.Module):
def __init__(self, in_channels, out_channels, dim=1, device='cuda'):
super().__init__()
# print('==> SharedMLP device: ', device)
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)
|