def __init__(self, in_channels_left, out_channels_left, in_channels_right, out_channels_right):
super(ReductionCell1, self).__init__()
self.conv_prev_1x1 = []
self.conv_prev_1x1.append(M.ReLU())
self.conv_prev_1x1.append(M.Conv2d(in_channels_left, out_channels_left, 1, stride=1, bias=False))
self.conv_prev_1x1.append(M.BatchNorm2d(out_channels_left, eps=0.001, momentum=0.1, affine=True))
self.conv_prev_1x1 = M.Sequential(*self.conv_prev_1x1)
self.conv_1x1 = []
self.conv_1x1.append(M.ReLU())
self.conv_1x1.append(M.Conv2d(in_channels_right, out_channels_right, 1, stride=1, bias=False))
self.conv_1x1.append(M.BatchNorm2d(out_channels_right, eps=0.001, momentum=0.1, affine=True))
self.conv_1x1 = M.Sequential(*self.conv_1x1)
self.comb_iter_0_left = BranchSeparables(out_channels_right, out_channels_right, 5, 2, 2, bias=False)
self.comb_iter_0_right = BranchSeparables(out_channels_right, out_channels_right, 7, 2, 3, bias=False)
self.comb_iter_1_left = M.MaxPool2d(3, stride=2, padding=1)
self.comb_iter_1_right = BranchSeparables(out_channels_right, out_channels_right, 7, 2, 3, bias=False)
self.comb_iter_2_left = M.AvgPool2d(3, stride=2, padding=1)
self.comb_iter_2_right = BranchSeparables(out_channels_right, out_channels_right, 5, 2, 2, bias=False)
self.comb_iter_3_right = M.AvgPool2d(3, stride=1, padding=1)
self.comb_iter_4_left = BranchSeparables(out_channels_right, out_channels_right, 3, 1, 1, bias=False)
self.comb_iter_4_right = M.MaxPool2d(3, stride=2, padding=1)
def __init__(self, in_ch=3, num_classes=1000):
'''
The AlexNet.
args:
in_ch: int, the number of channels of inputs
num_classes: int, the number of classes that need to predict
reference:
"One weird trick for parallelizing convolutional neural networks"<https://arxiv.org/abs/1404.5997>
'''
super(AlexNet, self).__init__()
#the part to extract feature
self.features = M.Sequential(
M.Conv2d(in_ch, 64, kernel_size=11, stride=4, padding=11 // 4),
M.ReLU(),
M.MaxPool2d(kernel_size=3, stride=2),
M.Conv2d(64, 192, kernel_size=5, padding=2),
M.ReLU(),
M.MaxPool2d(kernel_size=3, stride=2),
M.Conv2d(192, 384, kernel_size=3, stride=1, padding=1),
M.ReLU(),
M.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
M.ReLU(),
M.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
M.ReLU(),
M.MaxPool2d(kernel_size=3, stride=2),
)
#global avg pooling
self.avgpool = M.AdaptiveAvgPool2d((6, 6))
#classify part
self.classifier = M.Sequential(M.Dropout(),
M.Linear(256 * 6 * 6, 4096), M.ReLU(),
M.Dropout(), M.Linear(4096, 4096),
M.ReLU(), M.Linear(4096, num_classes))
def __init__(self):
super().__init__()
self.conv0 = M.Conv2d(1, 20, kernel_size=5, bias=False)
self.bn0 = M.BatchNorm2d(20)
self.relu0 = M.ReLU()
self.pool0 = M.MaxPool2d(2)
self.conv1 = M.Conv2d(20, 20, kernel_size=5, bias=False)
self.bn1 = M.BatchNorm2d(20)
self.relu1 = M.ReLU()
self.pool1 = M.MaxPool2d(2)
self.fc0 = M.Linear(500, 64, bias=True)
self.relu2 = M.ReLU()
self.fc1 = M.Linear(64, 10, bias=True)
def __init__(self, in_cha):
super(AlexNet, self).__init__()
self.conv1 = M.conv_bn.ConvBnRelu2d(in_cha,
24,
kernel_size=11,
stride=2,
padding=5)
self.pool1 = M.MaxPool2d(3, 2, 1)
self.conv2 = M.conv_bn.ConvBnRelu2d(24, 48, 5, 1, 2)
self.pool2 = M.MaxPool2d(3, 2, 1)
self.conv3 = M.conv_bn.ConvBnRelu2d(48, 96, 3, 1, 1)
self.conv4 = M.conv_bn.ConvBnRelu2d(96, 96, 3, 1, 1)
self.conv5 = M.conv_bn.ConvBn2d(96, 48, 3, 1, 1)
def __init__(self, in_cha, ch=48):
super(AlexNet_stride8, self).__init__()
assert ch % 2 == 0, "channel nums should % 2 = 0"
self.conv1 = M.conv_bn.ConvBnRelu2d(in_cha,
ch // 2,
kernel_size=11,
stride=2,
padding=5)
self.pool1 = M.MaxPool2d(3, 2, 1)
self.conv2 = M.conv_bn.ConvBnRelu2d(ch // 2, ch, 5, 1, 2)
self.pool2 = M.MaxPool2d(3, 2, 1)
self.conv3 = M.conv_bn.ConvBnRelu2d(ch, ch * 2, 3, 1, 1)
self.conv4 = M.conv_bn.ConvBnRelu2d(ch * 2, ch * 2, 3, 1, 1)
self.conv5 = M.conv_bn.ConvBn2d(ch * 2, ch, 3, 1, 1)
def __init__(self):
super().__init__()
self.classifier = None
if dist.get_rank() == 0:
self.features = M.Sequential(
M.ConvBn2d(3, 64, 7, stride=2, padding=3, bias=False),
M.MaxPool2d(kernel_size=3, stride=2, padding=1),
BasicBlock(64, 64, 1),
BasicBlock(64, 64, 1),
)
elif dist.get_rank() == 1:
self.features = M.Sequential(
BasicBlock(64, 128, 2),
BasicBlock(128, 128, 1),
)
elif dist.get_rank() == 2:
self.features = M.Sequential(
BasicBlock(128, 256, 2),
BasicBlock(256, 256, 1),
)
elif dist.get_rank() == 3:
self.features = M.Sequential(
BasicBlock(256, 512, 2),
BasicBlock(512, 512, 1),
)
self.classifier = M.Linear(512, 1000)
def __init__(self):
super().__init__()
# 单信道图片, 两层 5x5 卷积 + ReLU + 池化
self.conv1 = M.Conv2d(1, 6, 5)
self.relu1 = M.ReLU()
self.pool1 = M.MaxPool2d(2, 2)
self.conv2 = M.Conv2d(6, 16, 5)
self.relu2 = M.ReLU()
self.pool2 = M.MaxPool2d(2, 2)
# 两层全连接 + ReLU
self.fc1 = M.Linear(16 * 5 * 5, 120)
self.relu3 = M.ReLU()
self.fc2 = M.Linear(120, 84)
self.relu4 = M.ReLU()
# 分类器
self.classifier = M.Linear(84, 10)
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int,
max_pool=True,
max_pool_factor=1.0):
super(ConvBlock, self).__init__()
stride = (int(2 * max_pool_factor), int(2 * max_pool_factor))
if max_pool:
self.max_pool = M.MaxPool2d(kernel_size=stride, stride=stride)
stride = (1, 1)
else:
self.max_pool = lambda x: x
self.normalize = M.BatchNorm2d(out_channels, affine=True)
minit.uniform_(self.normalize.weight)
self.relu = M.ReLU()
self.conv = M.Conv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=1,
bias=True,
)
maml_init_(self.conv)
def __init__(self):
super(ResNet50, self).__init__()
self.conv1 = M.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=has_bias)
self.bn1 = FrozenBatchNorm2d(64)
self.maxpool = M.MaxPool2d(kernel_size=3, stride=2, padding=1)
block_counts = [3, 4, 6, 3]
bottleneck_channels_list = [64, 128, 256, 512]
out_channels_list = [256, 512, 1024, 2048]
stride_list = [1, 2, 2, 2]
in_channels = 64
self.layer1 = self._make_layer(block_counts[0], 64,
bottleneck_channels_list[0], out_channels_list[0], stride_list[0])
self.layer2 = self._make_layer(block_counts[1], out_channels_list[0],
bottleneck_channels_list[1], out_channels_list[1], stride_list[1])
self.layer3 = self._make_layer(block_counts[2], out_channels_list[1],
bottleneck_channels_list[2], out_channels_list[2], stride_list[2])
self.layer4 = self._make_layer(block_counts[3], out_channels_list[2],
bottleneck_channels_list[3], out_channels_list[3], stride_list[3])
for l in self.modules():
if isinstance(l, M.Conv2d):
M.init.msra_normal_(l.weight, mode="fan_in")
if has_bias:
M.init.fill_(l.bias, 0)
def __init__(self,
num_classes=1000,
aux_logits=True,
transform_input=False,
inception_blocks=None):
super(Inception3, self).__init__()
if inception_blocks is None:
inception_blocks = [
BasicConv2d, InceptionA, InceptionB, InceptionC, InceptionD,
InceptionE, InceptionAux
]
assert len(inception_blocks) == 7
conv_block = inception_blocks[0]
inception_a = inception_blocks[1]
inception_b = inception_blocks[2]
inception_c = inception_blocks[3]
inception_d = inception_blocks[4]
inception_e = inception_blocks[5]
inception_aux = inception_blocks[6]
self.aux_logits = aux_logits
self.transform_input = transform_input
self.Conv2d_1a_3x3 = conv_block(3, 32, kernel_size=3, stride=2)
self.Conv2d_2a_3x3 = conv_block(32, 32, kernel_size=3)
self.Conv2d_2b_3x3 = conv_block(32, 64, kernel_size=3, padding=1)
self.maxpool1 = M.MaxPool2d(kernel_size=3, stride=2)
self.Conv2d_3b_1x1 = conv_block(64, 80, kernel_size=1)
self.Conv2d_4a_3x3 = conv_block(80, 192, kernel_size=3)
self.maxpool2 = M.MaxPool2d(kernel_size=3, stride=2)
self.Mixed_5b = inception_a(192, pool_features=32)
self.Mixed_5c = inception_a(256, pool_features=64)
self.Mixed_5d = inception_a(288, pool_features=64)
self.Mixed_6a = inception_b(288)
self.Mixed_6b = inception_c(768, channels_7x7=128)
self.Mixed_6c = inception_c(768, channels_7x7=160)
self.Mixed_6d = inception_c(768, channels_7x7=160)
self.Mixed_6e = inception_c(768, channels_7x7=192)
if aux_logits:
self.AuxLogits = inception_aux(768, num_classes)
self.Mixed_7a = inception_d(768)
self.Mixed_7b = inception_e(1280)
self.Mixed_7c = inception_e(2048)
self.avgpool = M.AvgPool2d(8)
self.dropout = M.Dropout()
self.fc = M.Linear(2048, num_classes)
def __init__(self, in_channels, conv_block=None):
super(InceptionB, self).__init__()
if conv_block is None:
conv_block = BasicConv2d
self.branch3x3 = conv_block(in_channels, 384, kernel_size=3, stride=2)
self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1)
self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1)
self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, stride=2)
self.maxpool = M.MaxPool2d(3, 2)
def __init__(self,
channels,
residuals,
init_block_kernel_size,
init_block_channels,
maxpool_pad,
in_channels=3,
in_size=(224, 224),
num_classes=1000):
super(SqueezeNet, self).__init__()
self.in_size = in_size
self.num_classes = num_classes
self.feature = []
init_block = SqueezeInitBlock(in_channels=in_channels,
out_channels=init_block_channels,
kernel_size=init_block_kernel_size)
self.feature.append(init_block)
in_channels = init_block_channels
for i, channels_per_stage in enumerate(channels):
stage = []
pool = M.MaxPool2d(kernel_size=3, stride=2, padding=maxpool_pad[i])
stage.append(pool)
for j, out_channels in enumerate(channels_per_stage):
expand_channels = out_channels // 2
squeeze_channels = out_channels // 8
unit = FireUnit(in_channels=in_channels,
squeeze_channels=squeeze_channels,
expand1x1_channels=expand_channels,
expand3x3_channels=expand_channels,
residual=((residuals is not None)
and (residuals[i][j] == 1)))
stage.append(unit)
in_channels = out_channels
self.feature += stage
self.feature.append(M.Dropout(drop_prob=0.5))
self.feature = M.Sequential(*self.feature)
self.output = []
final_conv = M.Conv2d(in_channels=in_channels,
out_channels=num_classes,
kernel_size=1)
self.output.append(final_conv)
final_activ = M.ReLU()
self.output.append(final_activ)
final_pool = M.AvgPool2d(kernel_size=13, stride=1)
self.output.append(final_pool)
self.output = M.Sequential(*self.output)
def __init__(self, stem_filters, num_filters=42):
super(CellStem0, self).__init__()
self.num_filters = num_filters
self.stem_filters = stem_filters
self.conv_1x1 = []
self.conv_1x1.append(M.ReLU())
self.conv_1x1.append(M.Conv2d(self.stem_filters, self.num_filters, 1, stride=1, bias=False))
self.conv_1x1.append(M.BatchNorm2d(self.num_filters, eps=0.001, momentum=0.1, affine=True))
self.conv_1x1 = M.Sequential(*self.conv_1x1)
self.comb_iter_0_left = BranchSeparables(self.num_filters, self.num_filters, 5, 2, 2)
self.comb_iter_0_right = BranchSeparablesStem(self.stem_filters, self.num_filters, 7, 2, 3, bias=False)
self.comb_iter_1_left = M.MaxPool2d(3, stride=2, padding=1)
self.comb_iter_1_right = BranchSeparablesStem(self.stem_filters, self.num_filters, 7, 2, 3, bias=False)
self.comb_iter_2_left = M.AvgPool2d(3, stride=2, padding=1)
self.comb_iter_2_right = BranchSeparablesStem(self.stem_filters, self.num_filters, 5, 2, 2, bias=False)
self.comb_iter_3_right = M.AvgPool2d(3, stride=1, padding=1)
self.comb_iter_4_left = BranchSeparables(self.num_filters, self.num_filters, 3, 1, 1, bias=False)
self.comb_iter_4_right = M.MaxPool2d(3, stride=2, padding=1)
def __init__(self, input_size=112):
assert input_size == 112, f"expected input_size == 112, got {input_size}"
super().__init__()
self.input_size = input_size
self.stem = M.Sequential(
M.Conv2d(3, 8, kernel_size=3, stride=2, padding=1, bias=False),
M.BatchNorm2d(8),
M.ReLU(),
M.MaxPool2d(kernel_size=2, stride=2),
BasicBlock(8, 16),
BasicBlock(16, 32, stride=2),
BasicBlock(32, 64, stride=2),
)
self.fc = M.Linear(64, 9)
def __init__(self, stem_filters, num_filters):
super(CellStem1, self).__init__()
self.num_filters = num_filters
self.stem_filters = stem_filters
self.conv_1x1 = []
self.conv_1x1.append(M.ReLU())
self.conv_1x1.append(M.Conv2d(2*self.num_filters, self.num_filters, 1, stride=1, bias=False))
self.conv_1x1.append(M.BatchNorm2d(self.num_filters, eps=0.001, momentum=0.1, affine=True))
self.conv_1x1 = M.Sequential(*self.conv_1x1)
self.relu = M.ReLU()
self.path_1 = []
self.path_1.append(M.AvgPool2d(1, stride=2))
self.path_1.append(M.Conv2d(self.stem_filters, self.num_filters//2, 1, stride=1, bias=False))
self.path_1 = M.Sequential(*self.path_1)
self.path_2 = []
# self.path_2.add_module('pad', M.ZeroPad2d((0, 1, 0, 1)))
self.path_2.append(M.AvgPool2d(1, stride=2))
self.path_2.append(M.Conv2d(self.stem_filters, self.num_filters//2, 1, stride=1, bias=False))
self.path_2 = M.Sequential(*self.path_2)
self.final_path_bn = M.BatchNorm2d(self.num_filters, eps=0.001, momentum=0.1, affine=True)
self.comb_iter_0_left = BranchSeparables(self.num_filters, self.num_filters, 5, 2, 2, bias=False)
self.comb_iter_0_right = BranchSeparables(self.num_filters, self.num_filters, 7, 2, 3, bias=False)
self.comb_iter_1_left = M.MaxPool2d(3, stride=2, padding=1)
self.comb_iter_1_right = BranchSeparables(self.num_filters, self.num_filters, 7, 2, 3, bias=False)
self.comb_iter_2_left = M.AvgPool2d(3, stride=2, padding=1)
self.comb_iter_2_right = BranchSeparables(self.num_filters, self.num_filters, 5, 2, 2, bias=False)
self.comb_iter_3_right = M.AvgPool2d(3, stride=1, padding=1)
self.comb_iter_4_left = BranchSeparables(self.num_filters, self.num_filters, 3, 1, 1, bias=False)
self.comb_iter_4_right = M.MaxPool2d(3, stride=2, padding=1)
def __init__(self,
num_classes=1000,
aux_logits=True,
transform_input=False,
init_weights=True):
super(GoogLeNet, self).__init__()
self.aux_logits = aux_logits
self.transform_input = transform_input
self.conv1 = BasicConv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.maxpool1 = M.MaxPool2d(3, stride=2, padding=1) #向上取整
self.conv2 = BasicConv2d(64, 64, kernel_size=1)
self.conv3 = BasicConv2d(64, 192, kernel_size=3, padding=1)
self.maxpool2 = M.MaxPool2d(3, stride=2, padding=1)
self.inception3a = Inception(192, 64, 96, 128, 16, 32, 32)
self.inception3b = Inception(256, 128, 128, 192, 32, 96, 64)
self.maxpool3 = M.MaxPool2d(3, stride=2, padding=1)
self.inception4a = Inception(480, 192, 96, 208, 16, 48, 64)
self.inception4b = Inception(512, 160, 112, 224, 24, 64, 64)
self.inception4c = Inception(512, 128, 128, 256, 24, 64, 64)
self.inception4d = Inception(512, 112, 144, 288, 32, 64, 64)
self.inception4e = Inception(528, 256, 160, 320, 32, 128, 128)
self.maxpool4 = M.MaxPool2d(2, stride=2, padding=0)
self.inception5a = Inception(832, 256, 160, 320, 32, 128, 128)
self.inception5b = Inception(832, 384, 192, 384, 48, 128, 128)
if aux_logits:
self.aux1 = InceptionAux(512, num_classes)
self.aux2 = InceptionAux(528, num_classes)
# TODO
self.avgpool = M.AvgPool2d(7)
self.dropout = M.Dropout(0.2)
self.fc = M.Linear(1024, num_classes)
def __init__(self, in_channels, ch1x1, ch3x3red, ch3x3, ch5x5red, ch5x5,
pool_proj):
super(Inception, self).__init__()
self.branch1 = BasicConv2d(in_channels, ch1x1, kernel_size=1)
self.branch2 = M.Sequential(
BasicConv2d(in_channels, ch3x3red, kernel_size=1),
BasicConv2d(ch3x3red, ch3x3, kernel_size=3, padding=1))
self.branch3 = M.Sequential(
BasicConv2d(in_channels, ch5x5red, kernel_size=1),
BasicConv2d(ch5x5red, ch5x5, kernel_size=3, padding=1))
self.branch4 = M.Sequential(
M.MaxPool2d(kernel_size=3, stride=1, padding=1),
BasicConv2d(in_channels, pool_proj, kernel_size=1))
def __init__(self, in_ch, out_ch):
super(Unet, self).__init__()
self.conv1 = DoubleConv(in_ch, 32)
self.conv2 = DoubleConv(32, 64)
self.conv3 = DoubleConv(64, 128)
self.conv4 = DoubleConv(128, 256)
self.conv5 = DoubleConv(256, 512)
self.pool = M.MaxPool2d(2)
self.dropout = M.Dropout(0.3)
self.conv6 = DoubleConv(768, 256)
self.conv7 = DoubleConv(384, 128)
self.conv8 = DoubleConv(192, 64)
self.conv9 = DoubleConv(96, 32)
self.conv10 = M.Conv2d(32, out_ch, 1)
def __init__(self,
growth_rate=32,
block_config=(6, 12, 24, 16),
num_init_features=64,
bn_size=4,
drop_rate=0,
num_classes=1000):
super(DenseNet, self).__init__()
features = []
# First convolution
features.append(
M.Conv2d(3,
num_init_features,
kernel_size=7,
stride=2,
padding=3,
bias=False))
features.append(M.BatchNorm2d(num_init_features))
features.append(M.ReLU())
features.append(M.MaxPool2d(kernel_size=3, stride=2, padding=1))
# Each denseblock
num_features = num_init_features
for i, num_layers in enumerate(block_config):
block = _DenseBlock(num_layers=num_layers,
num_input_features=num_features,
bn_size=bn_size,
growth_rate=growth_rate,
drop_rate=drop_rate)
features.append(block)
num_features = num_features + num_layers * growth_rate
if i != len(block_config) - 1:
trans = _Transition(num_input_features=num_features,
num_output_features=num_features // 2)
features.append(trans)
num_features = num_features // 2
# Final batch norm
features.append(M.BatchNorm2d(num_features))
self.features = M.Sequential(*features)
# Linear layer
self.classifier = M.Linear(num_features, num_classes)
def __init__(self, in_channels, conv_block=None):
super(InceptionD, self).__init__()
if conv_block is None:
conv_block = BasicConv2d
self.branch3x3_1 = conv_block(in_channels, 192, kernel_size=1)
self.branch3x3_2 = conv_block(192, 320, kernel_size=3, stride=2)
self.branch7x7x3_1 = conv_block(in_channels, 192, kernel_size=1)
self.branch7x7x3_2 = conv_block(192,
192,
kernel_size=(1, 7),
padding=(0, 3))
self.branch7x7x3_3 = conv_block(192,
192,
kernel_size=(7, 1),
padding=(3, 0))
self.branch7x7x3_4 = conv_block(192, 192, kernel_size=3, stride=2)
self.maxpool = M.MaxPool2d(3, 2)
def _make_layers(self, in_channels, cfg, batch_norm=False):
'''
Make the layer from the config.
Args:
in_channels(int): the number of channels of first conv layer.
cfg(list): the config of the layers.
batch_norm(bool): If true use the batch normalization after the conv2d layer
'''
layers = []
in_ch = in_channels
for v in cfg:
if v == "M":
layers.append(M.MaxPool2d(kernel_size=2, stride=2))
else:
conv2d = M.Conv2d(in_ch, v, kernel_size=3, stride=1, padding=1)
if batch_norm:
layers += [conv2d, M.BatchNorm2d(v), M.ReLU()]
else:
layers += [conv2d, M.ReLU()]
in_ch = v
return M.Sequential(*layers)
def __init__(self,
in_channels,
out_channels,
kernel_sizes=(5, 9, 13),
activation="silu"):
super().__init__()
hidden_channels = in_channels // 2
self.conv1 = BaseConv(in_channels,
hidden_channels,
1,
stride=1,
act=activation)
self.m = [
M.MaxPool2d(kernel_size=ks, stride=1, padding=ks // 2)
for ks in kernel_sizes
]
conv2_channels = hidden_channels * (len(kernel_sizes) + 1)
self.conv2 = BaseConv(conv2_channels,
out_channels,
1,
stride=1,
act=activation)
def __init__(self, input_size=224, num_classes=1000, model_size="1.5x"):
super(ShuffleNetV2, self).__init__()
self.stage_repeats = [4, 8, 4]
self.model_size = model_size
if model_size == "0.5x":
self.stage_out_channels = [-1, 24, 48, 96, 192, 1024]
elif model_size == "1.0x":
self.stage_out_channels = [-1, 24, 116, 232, 464, 1024]
elif model_size == "1.5x":
self.stage_out_channels = [-1, 24, 176, 352, 704, 1024]
elif model_size == "2.0x":
self.stage_out_channels = [-1, 24, 244, 488, 976, 2048]
else:
raise NotImplementedError
# building first layer
input_channel = self.stage_out_channels[1]
self.first_conv = M.Sequential(
M.Conv2d(3, input_channel, 3, 2, 1, bias=False),
M.BatchNorm2d(input_channel),
FReLU(input_channel),
)
self.maxpool = M.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.features = []
for idxstage in range(len(self.stage_repeats)):
numrepeat = self.stage_repeats[idxstage]
output_channel = self.stage_out_channels[idxstage + 2]
for i in range(numrepeat):
if i == 0:
self.features.append(
ShuffleV2Block(
input_channel,
output_channel,
mid_channels=output_channel // 2,
ksize=3,
stride=2,
))
else:
self.features.append(
ShuffleV2Block(
input_channel // 2,
output_channel,
mid_channels=output_channel // 2,
ksize=3,
stride=1,
))
input_channel = output_channel
self.features = M.Sequential(*self.features)
self.conv_last = M.Sequential(
M.Conv2d(input_channel,
self.stage_out_channels[-1],
1,
1,
0,
bias=False),
M.BatchNorm2d(self.stage_out_channels[-1]),
FReLU(self.stage_out_channels[-1]),
)
self.globalpool = M.AvgPool2d(7)
if self.model_size == "2.0x":
self.dropout = M.Dropout(0.2)
self.classifier = M.Sequential(
M.Linear(self.stage_out_channels[-1], num_classes, bias=False))
self._initialize_weights()
def __init__(self, in_feature="p5"):
super().__init__()
self.num_levels = 1
self.in_feature = in_feature
self.pool = M.MaxPool2d(kernel_size=1, stride=2, padding=0)
def __init__(
self,
block,
layers=[2, 2, 2, 2],
num_classes=1000,
zero_init_residual=False,
groups=1,
width_per_group=64,
replace_stride_with_dilation=None,
norm=M.BatchNorm2d,
):
super().__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.conv1 = M.Conv2d(3,
self.in_channels,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = norm(self.in_channels)
self.maxpool = M.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1_0 = BasicBlock(
self.in_channels,
64,
stride=1,
groups=self.groups,
base_width=self.base_width,
dilation=self.dilation,
norm=M.BatchNorm2d,
)
self.layer1_1 = BasicBlock(
self.in_channels,
64,
stride=1,
groups=self.groups,
base_width=self.base_width,
dilation=self.dilation,
norm=M.BatchNorm2d,
)
self.layer2_0 = BasicBlock(64, 128, stride=2)
self.layer2_1 = BasicBlock(128, 128)
self.layer3_0 = BasicBlock(128, 256, stride=2)
self.layer3_1 = BasicBlock(256, 256)
self.layer4_0 = BasicBlock(256, 512, stride=2)
self.layer4_1 = BasicBlock(512, 512)
self.layer1 = self._make_layer(block, 64, layers[0], norm=norm)
self.layer2 = self._make_layer(block,
128,
2,
stride=2,
dilate=replace_stride_with_dilation[0],
norm=norm)
self.layer3 = self._make_layer(block,
256,
2,
stride=2,
dilate=replace_stride_with_dilation[1],
norm=norm)
self.layer4 = self._make_layer(block,
512,
2,
stride=2,
dilate=replace_stride_with_dilation[2],
norm=norm)
self.fc = M.Linear(512, 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)
if zero_init_residual:
for m in self.modules():
M.init.zeros_(m.bn2.weight)
def __init__(self,
block,
blocks,
in_ch=3,
num_classes=1000,
first_stride=2,
light_head=False,
zero_init_residual=False,
groups=1,
width_per_group=64,
strides=[1, 2, 2, 2],
dilations=[1, 1, 1, 1],
multi_grids=[1, 1, 1],
norm_layer=None,
se_module=None,
reduction=16,
radix=0,
avd=False,
avd_first=False,
avg_layer=False,
avg_down=False,
stem_width=64):
'''
Modified resnet according to https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
Implementate ResNet and the variation of ResNet.
Args:
in_ch: int, the number of channels of the input
block: BasicBlock or Bottleneck.The block of the resnet
num_classes: int, the number of classes to predict
first_stride: int, the stride of the first conv layer
light_head: boolean, whether use conv3x3 replace the conv7x7 in first conv layer
zero_init_residual: whether initilize the residule block's batchnorm with zero
groups: int, the number of groups for the conv in net
width_per_group: int, the width of the conv layers
strides: list, the list of the strides for the each stage
dilations: list, the dilations of each block
multi_grids: list, implementation of the multi grid layer in deeplabv3
norm_layer: megengine.module.Module, the normalization layer, default is batch normalization
se_module: SEModule, the Squeeze Excitation Module
radix: int, the radix index from ResNest
reduction: int, the reduction rate
avd: bool, whether use the avd layer
avd_first: bool, whether use the avd layer before bottleblock's conv2
stem_width: int, the channels of the conv3x3 when use 3 conv3x3 replace conv7x7
References:
"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>
"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>
https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch
deeplab v3: https://arxiv.org/pdf/1706.05587.pdf
deeplab v3+: https://arxiv.org/pdf/1802.02611.pdf
"Squeeze-and-Excitation Networks"<https://arxiv.org/abs/1709.01507>
"ResNeSt: Split-Attention Networks"<https://arxiv.org/pdf/2004.08955.pdf>
'''
super(ResNet, self).__init__()
if len(dilations) != 4:
raise ValueError(
"The length of dilations must be 4, but got {}".format(
len(dilations)))
if len(strides) != 4:
raise ValueError(
"The length of dilations must be 4, but got {}".format(
len(strides)))
if len(multi_grids) > blocks[-1]:
multi_grids = multi_grids[:blocks[-1]]
elif len(multi_grids) < blocks[-1]:
raise ValueError(
"The length of multi_grids must greater than or equal the number of blocks for last stage , but got {}/{}"
.format(len(multi_grids), blocks[-1]))
if norm_layer is None:
norm_layer = M.BatchNorm2d
self.base_width = width_per_group
self.multi_grids = multi_grids
self.inplanes = 64
self.groups = groups
self.norm_layer = norm_layer
self.avg_layer = avg_layer
self.avg_down = avg_down
if light_head:
self.conv1 = M.Sequential(
conv3x3(in_ch, stem_width, stride=first_stride),
norm_layer(stem_width),
M.ReLU(),
conv3x3(stem_width, stem_width, stride=1),
norm_layer(stem_width),
M.ReLU(),
conv3x3(stem_width, self.inplanes, stride=1),
)
else:
self.conv1 = M.Conv2d(in_ch,
self.inplanes,
kernel_size=7,
stride=first_stride,
padding=3,
bias=False)
self.bn1 = norm_layer(self.inplanes)
self.relu = M.ReLU()
self.maxpool = M.MaxPool2d(kernel_size=3, stride=2, padding=1)
#4 stage
self.layer1 = self._make_layer(block,
64,
blocks[0],
stride=strides[0],
dilation=dilations[0],
se_module=se_module,
reduction=reduction,
radix=radix,
avd=avd,
avd_first=avd_first)
self.layer2 = self._make_layer(block,
128,
blocks[1],
stride=strides[1],
dilation=dilations[1],
se_module=se_module,
reduction=reduction,
radix=radix,
avd=avd,
avd_first=avd_first)
self.layer3 = self._make_layer(block,
256,
blocks[2],
stride=strides[2],
dilation=dilations[2],
se_module=se_module,
reduction=reduction,
radix=radix,
avd=avd,
avd_first=avd_first)
self.layer4 = self._make_grid_layer(block,
512,
blocks[3],
stride=strides[3],
dilation=dilations[3],
se_module=se_module,
reduction=reduction,
radix=radix,
avd=avd,
avd_first=avd_first)
#classification part
self.avgpool = M.AdaptiveAvgPool2d(1)
self.fc = M.Linear(self.inplanes, num_classes)
for m in self.modules():
if isinstance(m, M.Conv2d):
M.init.msra_normal_(m.weight,
mode="fan_out",
nonlinearity="relu")
elif isinstance(m, M.BatchNorm2d):
M.init.fill_(m.weight, 1)
M.init.zeros_(m.bias)
# 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)
def __init__(
self,
layers,
num_classes=10,
zero_init_residual=False,
groups=1,
width_per_group=64,
replace_stride_with_dilation=None,
norm=M.BatchNorm2d,
):
block = Bottleneck
super(ResNet_MNIST, 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.conv1 = M.Conv2d(
#3, self.in_channels, kernel_size=7, stride=2, padding=3, bias=False
1,
self.in_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.bn1 = norm(self.in_channels)
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.fc = M.Linear(512 * block.expansion, num_classes)
self.fc = M.Linear(512, 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)