def __init__(
self,
in_channels,
channels,
stride=1,
groups=1,
base_width=64,
dilation=1,
norm=M.BatchNorm2d,
):
assert norm is M.BatchNorm2d, "Quant mode only support BatchNorm2d currently."
super(Bottleneck, self).__init__()
width = int(channels * (base_width / 64.0)) * groups
self.conv_bn_relu1 = M.ConvBnRelu2d(in_channels, width, 1, 1, bias=False)
self.conv_bn_relu2 = M.ConvBnRelu2d(
width,
width,
3,
stride,
padding=dilation,
groups=groups,
dilation=dilation,
bias=False,
)
self.conv_bn3 = M.ConvBn2d(width, channels * self.expansion, 1, 1, bias=False)
self.downsample = (
M.Identity()
if in_channels == channels * self.expansion and stride == 1
else M.ConvBn2d(
in_channels, channels * self.expansion, 1, stride, bias=False
)
)
self.add = M.Elemwise("FUSE_ADD_RELU")
def __init__(self, inp, oup, stride, expand_ratio):
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(
M.ConvBnRelu2d(inp, hidden_dim, kernel_size=1, bias=False))
layers.extend([
# dw
M.ConvBnRelu2d(
hidden_dim,
hidden_dim,
kernel_size=3,
padding=1,
stride=stride,
groups=hidden_dim,
bias=False,
),
# pw-linear
M.ConvBn2d(hidden_dim, oup, kernel_size=1, bias=False),
])
self.conv = M.Sequential(*layers)
self.add = M.Elemwise("ADD")
def __init__(
self,
in_channels,
channels,
stride=1,
groups=1,
base_width=64,
dilation=1,
norm=M.BatchNorm2d,
):
assert norm is M.BatchNorm2d, "Quant mode only support BatchNorm2d currently."
super(BasicBlock, self).__init__()
if groups != 1 or base_width != 64:
raise ValueError("BasicBlock only supports groups=1 and base_width=64")
if dilation > 1:
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
self.conv_bn_relu1 = M.ConvBnRelu2d(
in_channels, channels, 3, stride, padding=dilation, bias=False
)
self.conv_bn2 = M.ConvBn2d(channels, channels, 3, 1, padding=1, bias=False)
self.downsample = (
M.Identity()
if in_channels == channels and stride == 1
else M.ConvBn2d(in_channels, channels, 1, stride, bias=False)
)
self.add = M.Elemwise("FUSE_ADD_RELU")
def __init__(self, has_bias=True):
super().__init__()
self.data = np.random.random((1, 3, 224, 224)).astype(np.float32)
self.convbnrelu = M.ConvBnRelu2d(3,
10,
3,
stride=(2, 3),
dilation=(2, 2),
padding=(3, 1),
bias=has_bias)
def __init__(self, inp, oup, stride):
super().__init__()
if inp == oup and stride == 1:
self.proj = M.Identity()
else:
self.proj = M.ConvBn2d(inp, oup, 1, stride=stride, bias=False)
self.conv1 = M.ConvBnRelu2d(inp,
oup,
3,
padding=1,
stride=stride,
bias=False)
self.conv2 = M.ConvBn2d(oup, oup, 3, padding=1, stride=1, bias=False)
def _init_reg_convs(self):
"""Initialize bbox regression conv layers of the head."""
self.reg_convs = []
for i in range(self.stacked_convs - 1):
self.reg_convs.append(
M.ConvRelu2d(self.feat_channels, self.feat_channels, 3, 1, 1))
self.reg_convs.append(
M.ConvBnRelu2d(self.feat_channels,
self.feat_channels,
kernel_size=3,
stride=1,
padding=1,
bias=True,
momentum=0.9,
affine=True,
track_running_stats=True))
self.reg_convs = M.Sequential(*self.reg_convs)
def __init__(self):
super().__init__()
self.conv = Float.ConvBnRelu2d(3, 3, 3)
self.conv.disable_quantize()
def __init__(self):
super().__init__()
self.conv_bn = M.ConvBnRelu2d(3, 3, 3)
self.linear = M.Linear(3, 3)
def __init__(
self,
num_classes=1000,
width_mult=1.0,
inverted_residual_setting=None,
round_nearest=8,
):
"""
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
"""
super(MobileNetV2, self).__init__()
block = InvertedResidual
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 = [
M.ConvBnRelu2d(3,
input_channel,
kernel_size=3,
padding=1,
stride=2,
bias=False)
]
# 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))
input_channel = output_channel
# building last several layers
features.append(
M.ConvBnRelu2d(input_channel,
self.last_channel,
kernel_size=1,
bias=False))
# make it M.Sequential
self.features = M.Sequential(*features)
# building classifier
self.classifier = M.Sequential(
M.Dropout(0.2),
M.Linear(self.last_channel, num_classes),
)
self.classifier.disable_quantize()
self.quant = M.QuantStub()
self.dequant = M.DequantStub()
# weight initialization
for m in self.modules():
if isinstance(m, M.Conv2d):
M.init.msra_normal_(m.weight, mode="fan_out")
if m.bias is not None:
M.init.zeros_(m.bias)
elif isinstance(m, M.BatchNorm2d):
M.init.ones_(m.weight)
M.init.zeros_(m.bias)
elif isinstance(m, M.Linear):
M.init.normal_(m.weight, 0, 0.01)
M.init.zeros_(m.bias)
def __init__(
self,
block,
layers,
num_classes=1000,
zero_init_residual=False,
groups=1,
width_per_group=64,
replace_stride_with_dilation=None,
norm=M.BatchNorm2d,
):
super(ResNet, self).__init__()
self.in_channels = 64
self.dilation = 1
if replace_stride_with_dilation is None:
# each element in the tuple indicates if we should replace
# the 2x2 stride with a dilated convolution instead
replace_stride_with_dilation = [False, False, False]
if len(replace_stride_with_dilation) != 3:
raise ValueError(
"replace_stride_with_dilation should be None "
"or a 3-element tuple, got {}".format(replace_stride_with_dilation)
)
self.groups = groups
self.base_width = width_per_group
self.quant = M.QuantStub()
self.dequant = M.DequantStub()
self.conv_bn_relu1 = M.ConvBnRelu2d(
3, self.in_channels, kernel_size=7, stride=2, padding=3, bias=False
)
self.maxpool = M.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0], norm=norm)
self.layer2 = self._make_layer(
block,
128,
layers[1],
stride=2,
dilate=replace_stride_with_dilation[0],
norm=norm,
)
self.layer3 = self._make_layer(
block,
256,
layers[2],
stride=2,
dilate=replace_stride_with_dilation[1],
norm=norm,
)
self.layer4 = self._make_layer(
block,
512,
layers[3],
stride=2,
dilate=replace_stride_with_dilation[2],
norm=norm,
)
self.fc = M.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, M.Conv2d):
M.init.msra_normal_(m.weight, mode="fan_out", nonlinearity="relu")
if m.bias is not None:
fan_in, _ = M.init.calculate_fan_in_and_fan_out(m.weight)
bound = 1 / math.sqrt(fan_in)
M.init.uniform_(m.bias, -bound, bound)
elif isinstance(m, M.BatchNorm2d):
M.init.ones_(m.weight)
M.init.zeros_(m.bias)
elif isinstance(m, M.Linear):
M.init.msra_uniform_(m.weight, a=math.sqrt(5))
if m.bias is not None:
fan_in, _ = M.init.calculate_fan_in_and_fan_out(m.weight)
bound = 1 / math.sqrt(fan_in)
M.init.uniform_(m.bias, -bound, bound)
# Zero-initialize the last BN in each residual branch,
# so that the residual branch starts with zeros,
# and each residual block behaves like an identity.
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
if zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
M.init.zeros_(m.bn3.weight)
elif isinstance(m, BasicBlock):
M.init.zeros_(m.bn2.weight)
