def test_aspp():
# test aspp with normal conv
aspp = ASPP(128, out_channels=512, mid_channels=128, dilations=(6, 12, 18))
assert aspp.convs[0].activate.__class__.__name__ == 'ReLU'
assert aspp.convs[0].conv.out_channels == 128
assert aspp.convs[1].__class__.__name__ == 'ConvModule'
for conv_idx in range(1, 4):
assert aspp.convs[conv_idx].conv.dilation[0] == 6 * conv_idx
x = torch.rand(2, 128, 8, 8)
output = aspp(x)
assert output.shape == (2, 512, 8, 8)
# test aspp with separable conv
aspp = ASPP(128, separable_conv=True)
assert aspp.convs[1].__class__.__name__ == 'DepthwiseSeparableConvModule'
x = torch.rand(2, 128, 8, 8)
output = aspp(x)
assert output.shape == (2, 256, 8, 8)
# test aspp with ReLU6
aspp = ASPP(128, dilations=(12, 24, 36), act_cfg=dict(type='ReLU6'))
assert aspp.convs[0].activate.__class__.__name__ == 'ReLU6'
for conv_idx in range(1, 4):
assert aspp.convs[conv_idx].conv.dilation[0] == 12 * conv_idx
x = torch.rand(2, 128, 8, 8)
output = aspp(x)
assert output.shape == (2, 256, 8, 8)
def __init__(self,
in_channels,
batch_norm=False,
aspp=False,
dilations=None):
super(VGG16, self).__init__()
self.batch_norm = batch_norm
self.aspp = aspp
self.dilations = dilations
self.layer1 = self._make_layer(in_channels, 64, 2)
self.layer2 = self._make_layer(64, 128, 2)
self.layer3 = self._make_layer(128, 256, 3)
self.layer4 = self._make_layer(256, 512, 3)
self.layer5 = self._make_layer(512, 512, 3)
self.conv6 = nn.Conv2d(512, 512, kernel_size=1)
if self.batch_norm:
self.bn = nn.BatchNorm2d(512)
self.relu = nn.ReLU(inplace=True)
if self.aspp:
self.aspp = ASPP(512, dilations=self.dilations)
self.out_channels = 256
else:
self.out_channels = 512
def __init__(self,
in_channels,
out_stride=32,
width_mult=1,
index_mode='m2o',
aspp=True,
norm_cfg=dict(type='BN'),
freeze_bn=False,
use_nonlinear=True,
use_context=True):
super().__init__()
if out_stride not in [16, 32]:
raise ValueError(f'out_stride must 16 or 32, got {out_stride}')
self.out_stride = out_stride
self.width_mult = width_mult
# we name the index network in the paper index_block
if index_mode == 'holistic':
index_block = HolisticIndexBlock
elif index_mode in ('o2o', 'm2o'):
index_block = partial(DepthwiseIndexBlock, mode=index_mode)
else:
raise NameError('Unknown index block mode {}'.format(index_mode))
# default setting
initial_channels = 32
inverted_residual_setting = [
# expand_ratio, input_chn, output_chn, num_blocks, stride, dilation
[1, initial_channels, 16, 1, 1, 1],
[6, 16, 24, 2, 2, 1],
[6, 24, 32, 3, 2, 1],
[6, 32, 64, 4, 2, 1],
[6, 64, 96, 3, 1, 1],
[6, 96, 160, 3, 2, 1],
[6, 160, 320, 1, 1, 1],
]
# update layer setting according to width_mult
initial_channels = int(initial_channels * width_mult)
for layer_setting in inverted_residual_setting:
# update in_channels and out_channels
layer_setting[1] = int(layer_setting[1] * self.width_mult)
layer_setting[2] = int(layer_setting[2] * self.width_mult)
if out_stride == 32:
# It should be noted that layers 0 is not an InvertedResidual layer
# but a ConvModule. Thus, the index of InvertedResidual layer in
# downsampled_layers starts at 1.
self.downsampled_layers = [0, 2, 3, 4, 6]
else: # out_stride is 16
self.downsampled_layers = [0, 2, 3, 4]
# if out_stride is 16, then increase the dilation of the last two
# InvertedResidual layer to increase the receptive field
inverted_residual_setting[5][5] = 2
inverted_residual_setting[6][5] = 2
# build the first layer
self.layers = nn.ModuleList([
ConvModule(in_channels,
initial_channels,
3,
padding=1,
norm_cfg=norm_cfg,
act_cfg=dict(type='ReLU6'))
])
# build bottleneck layers
for layer_setting in inverted_residual_setting:
self.layers.append(self._make_layer(layer_setting, norm_cfg))
# freeze encoder batch norm layers
if freeze_bn:
self.freeze_bn()
# build index blocks
self.index_layers = nn.ModuleList()
for layer in self.downsampled_layers:
# inverted_residual_setting begins at layer1, the in_channels
# of layer1 is the out_channels of layer0
self.index_layers.append(
index_block(inverted_residual_setting[layer][1], norm_cfg,
use_context, use_nonlinear))
self.avg_pool = nn.AvgPool2d(2, stride=2)
if aspp:
dilation = (2, 4, 8) if out_stride == 32 else (6, 12, 18)
self.dconv = ASPP(320 * self.width_mult,
160,
mid_channels=int(256 * self.width_mult),
dilations=dilation,
norm_cfg=norm_cfg,
act_cfg=dict(type='ReLU6'),
separable_conv=True)
else:
self.dconv = ConvModule(320 * self.width_mult,
160,
1,
norm_cfg=norm_cfg,
act_cfg=dict(type='ReLU6'))
self.out_channels = 160