Spaces:
Sleeping
Sleeping
File size: 1,032 Bytes
2fd6166 |
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 |
import torch.nn as nn
import torch
__all__ = ['SharedMLP']
class Swish(nn.Module):
def forward(self,x):
return x * torch.sigmoid(x)
class SharedMLP(nn.Module):
def __init__(self, in_channels, out_channels, dim=1):
super().__init__()
if dim == 1: # default value
conv = nn.Conv1d
bn = nn.GroupNorm
elif dim == 2:
conv = nn.Conv2d
bn = nn.GroupNorm
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),
bn(8, oc),
Swish(),
])
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)
|