def __init__(self, inp, oup, stride, expand_ratio, norm_layer):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
layers = []
if expand_ratio != 1:
# pw
layers.append(
ConvBNReLU(inp,
hidden_dim,
kernel_size=1,
norm_layer=norm_layer))
layers.extend([
# dw
ConvBNReLU(
hidden_dim,
hidden_dim,
stride=stride,
groups=hidden_dim,
norm_layer=norm_layer,
),
# pw-linear
Lift(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False)),
Lift(norm_layer(oup)),
])
self.conv = nn.Sequential(*layers)
def make_layers(cfg, batch_norm=False):
layers = []
in_channels = 3
for v in cfg:
if v == "M":
layers += [Lift(nn.MaxPool2d(kernel_size=2, stride=2))]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [Lift(conv2d), Lift(nn.BatchNorm2d(v)), LIF()]
else:
layers += [Lift(conv2d), LIF()]
in_channels = v
return nn.Sequential(*layers)
def test_lift_conv():
batch_size = 16
seq_length = 20
in_channels = 64
out_channels = 32
conv2d = Lift(torch.nn.Conv2d(in_channels, out_channels, 5, 1))
data = torch.randn(seq_length, batch_size, in_channels, 20, 30)
output = conv2d(data)
assert output.shape == torch.Size([seq_length, batch_size, out_channels, 16, 26])
def __init__(
self, in_planes, out_planes, norm_layer, kernel_size=3, stride=1, groups=1
):
padding = (kernel_size - 1) // 2
super(ConvBNReLU, self).__init__(
Lift(
nn.Conv2d(
in_planes,
out_planes,
kernel_size,
stride,
padding,
groups=groups,
bias=False,
)
),
Lift(norm_layer(out_planes)),
LIFFeedForwardLayer(),
)
def test_lift_sequential():
batch_size = 16
seq_length = 20
in_channels = 64
out_channels = 32
data = torch.randn(seq_length, batch_size, in_channels, 20, 30)
module = torch.nn.Sequential(
Lift(torch.nn.Conv2d(in_channels, out_channels, 5, 1)),
LIFFeedForwardLayer(),
)
output, _ = module(data)
assert output.shape == torch.Size([seq_length, batch_size, out_channels, 16, 26])
def __init__(self, features, num_classes=1000, init_weights=True):
super(VGG, self).__init__()
self.features = features
self.avgpool = Lift(nn.AdaptiveAvgPool2d((7, 7)))
self.classifier = nn.Sequential(
Lift(nn.Linear(512 * 7 * 7, 4096)),
LIFFeedForwardLayer(),
Lift(nn.Dropout()),
Lift(nn.Linear(4096, 4096)),
LIFFeedForwardLayer(),
Lift(nn.Dropout()),
Lift(nn.Linear(4096, num_classes)),
)
if init_weights:
self._initialize_weights()
def __init__(
self,
num_classes=1000,
width_mult=1.0,
inverted_residual_setting=None,
round_nearest=8,
block=None,
norm_layer=None,
):
"""
MobileNet V2 main class
Args:
num_classes (int): Number of classes
width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
inverted_residual_setting: Network structure
round_nearest (int): Round the number of channels in each layer to be a multiple of this number
Set to 1 to turn off rounding
block: Module specifying inverted residual building block for mobilenet
norm_layer: Module specifying the normalization layer to use
"""
super(MobileNetV2, self).__init__()
if block is None:
block = InvertedResidual
if norm_layer is None:
norm_layer = nn.BatchNorm2d
input_channel = 32
last_channel = 1280
if inverted_residual_setting is None:
inverted_residual_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
# only check the first element, assuming user knows t,c,n,s are required
if (len(inverted_residual_setting) == 0
or len(inverted_residual_setting[0]) != 4):
raise ValueError("inverted_residual_setting should be non-empty "
"or a 4-element list, got {}".format(
inverted_residual_setting))
# building first layer
input_channel = _make_divisible(input_channel * width_mult,
round_nearest)
self.last_channel = _make_divisible(
last_channel * max(1.0, width_mult), round_nearest)
features = [
ConvBNReLU(3, input_channel, stride=2, norm_layer=norm_layer)
]
# building inverted residual blocks
for t, c, n, s in inverted_residual_setting:
output_channel = _make_divisible(c * width_mult, round_nearest)
for i in range(n):
stride = s if i == 0 else 1
features.append(
block(
input_channel,
output_channel,
stride,
expand_ratio=t,
norm_layer=norm_layer,
))
input_channel = output_channel
# building last several layers
features.append(
ConvBNReLU(input_channel,
self.last_channel,
kernel_size=1,
norm_layer=norm_layer))
# make it nn.Sequential
self.features = nn.Sequential(*features)
self.pool = Lift(nn.AdaptiveAvgPool2d(1))
# building classifier
self.classifier = nn.Sequential(
Lift(nn.Dropout(0.2)),
Lift(nn.Linear(self.last_channel, num_classes)),
)
# weight initialization
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode="fan_out")
assert m.bias is None
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.zeros_(m.bias)