Fix the mid_channels of SEResNeXt

This commit is contained in:
yanglei 2020-07-07 11:25:48 +08:00 committed by chenkai
parent 9b0df9ea12
commit 17092d8be4
2 changed files with 37 additions and 12 deletions

View File

@ -12,6 +12,7 @@ class SEBottleneck(_SEBottleneck):
Args:
in_channels (int): Input channels of this block.
out_channels (int): Output channels of this block.
base_channels (int): Middle channels of the first stage. Default: 64.
groups (int): Groups of conv2.
width_per_group (int): Width per group of conv2. 64x4d indicates
``groups=64, width_per_group=4`` and 32x8d indicates
@ -35,6 +36,7 @@ class SEBottleneck(_SEBottleneck):
def __init__(self,
in_channels,
out_channels,
base_channels=64,
groups=32,
width_per_group=4,
se_ratio=16,
@ -44,30 +46,34 @@ class SEBottleneck(_SEBottleneck):
self.groups = groups
self.width_per_group = width_per_group
if groups == 1:
width = self.mid_channels
else:
width = groups * width_per_group
# We follow the same rational of ResNext to compute mid_channels.
# For SEResNet bottleneck, middle channels are determined by expansion
# and out_channels, but for SEResNeXt bottleneck, it is determined by
# groups and width_per_group and the stage it is located in.
if groups != 1:
assert self.mid_channels % base_channels == 0
self.mid_channels = (
groups * width_per_group * self.mid_channels // base_channels)
self.norm1_name, norm1 = build_norm_layer(
self.norm_cfg, width, postfix=1)
self.norm_cfg, self.mid_channels, postfix=1)
self.norm2_name, norm2 = build_norm_layer(
self.norm_cfg, width, postfix=2)
self.norm_cfg, self.mid_channels, postfix=2)
self.norm3_name, norm3 = build_norm_layer(
self.norm_cfg, self.out_channels, postfix=3)
self.conv1 = build_conv_layer(
self.conv_cfg,
self.in_channels,
width,
self.mid_channels,
kernel_size=1,
stride=self.conv1_stride,
bias=False)
self.add_module(self.norm1_name, norm1)
self.conv2 = build_conv_layer(
self.conv_cfg,
width,
width,
self.mid_channels,
self.mid_channels,
kernel_size=3,
stride=self.conv2_stride,
padding=self.dilation,
@ -77,7 +83,11 @@ class SEBottleneck(_SEBottleneck):
self.add_module(self.norm2_name, norm2)
self.conv3 = build_conv_layer(
self.conv_cfg, width, self.out_channels, kernel_size=1, bias=False)
self.conv_cfg,
self.mid_channels,
self.out_channels,
kernel_size=1,
bias=False)
self.add_module(self.norm3_name, norm3)
@ -138,4 +148,7 @@ class SEResNeXt(SEResNet):
def make_res_layer(self, **kwargs):
return ResLayer(
groups=self.groups, width_per_group=self.width_per_group, **kwargs)
groups=self.groups,
width_per_group=self.width_per_group,
base_channels=self.base_channels,
**kwargs)

View File

@ -13,12 +13,24 @@ def test_bottleneck():
# Test SEResNeXt Bottleneck structure
block = SEBottleneckX(
64, 256, groups=32, width_per_group=4, stride=2, style='pytorch')
assert block.width_per_group == 4
assert block.conv2.stride == (2, 2)
assert block.conv2.groups == 32
assert block.conv2.out_channels == 128
assert block.conv2.out_channels == block.mid_channels
# Test SEResNeXt Bottleneck structure (groups=1)
block = SEBottleneckX(
64, 256, groups=1, width_per_group=4, stride=2, style='pytorch')
assert block.conv2.stride == (2, 2)
assert block.conv2.groups == 1
assert block.conv2.out_channels == 64
assert block.mid_channels == 64
assert block.conv2.out_channels == block.mid_channels
# Test SEResNeXt Bottleneck forward
block = SEBottleneckX(64, 64, groups=32, width_per_group=4)
block = SEBottleneckX(
64, 64, base_channels=16, groups=32, width_per_group=4)
x = torch.randn(1, 64, 56, 56)
x_out = block(x)
assert x_out.shape == torch.Size([1, 64, 56, 56])