def test_res_layer():
# Test ResLayer of 3 Bottleneck w\o downsample
layer = ResLayer(SEBottleneck, 3, 64, 64, se_ratio=16)
assert len(layer) == 3
assert layer[0].conv1.in_channels == 64
assert layer[0].conv1.out_channels == 16
for i in range(1, len(layer)):
assert layer[i].conv1.in_channels == 64
assert layer[i].conv1.out_channels == 16
for i in range(len(layer)):
assert layer[i].downsample is None
x = torch.randn(1, 64, 56, 56)
x_out = layer(x)
assert x_out.shape == torch.Size([1, 64, 56, 56])
# Test ResLayer of 3 SEBottleneck with downsample
layer = ResLayer(SEBottleneck, 3, 64, 256, se_ratio=16)
assert layer[0].downsample[0].out_channels == 256
for i in range(1, len(layer)):
assert layer[i].downsample is None
x = torch.randn(1, 64, 56, 56)
x_out = layer(x)
assert x_out.shape == torch.Size([1, 256, 56, 56])
# Test ResLayer of 3 SEBottleneck with stride=2
layer = ResLayer(SEBottleneck, 3, 64, 256, stride=2, se_ratio=8)
assert layer[0].downsample[0].out_channels == 256
assert layer[0].downsample[0].stride == (2, 2)
for i in range(1, len(layer)):
assert layer[i].downsample is None
x = torch.randn(1, 64, 56, 56)
x_out = layer(x)
assert x_out.shape == torch.Size([1, 256, 28, 28])
# Test ResLayer of 3 SEBottleneck with stride=2 and average downsample
layer = ResLayer(SEBottleneck,
3,
64,
256,
stride=2,
avg_down=True,
se_ratio=8)
assert isinstance(layer[0].downsample[0], AvgPool2d)
assert layer[0].downsample[1].out_channels == 256
assert layer[0].downsample[1].stride == (1, 1)
for i in range(1, len(layer)):
assert layer[i].downsample is None
x = torch.randn(1, 64, 56, 56)
x_out = layer(x)
assert x_out.shape == torch.Size([1, 256, 28, 28])
def test_bottleneck_reslayer():
# 3 Bottleneck w/o downsample
layer = ResLayer(Bottleneck, 3, 32, 32)
assert len(layer) == 3
for i in range(3):
assert layer[i].in_channels == 32
assert layer[i].out_channels == 32
assert layer[i].downsample is None
x = torch.randn(1, 32, 56, 56)
x_out = layer(x)
assert x_out.shape == (1, 32, 56, 56)
# 3 Bottleneck w/ stride 1 and downsample
layer = ResLayer(Bottleneck, 3, 32, 64)
assert len(layer) == 3
assert layer[0].in_channels == 32
assert layer[0].out_channels == 64
assert layer[0].stride == 1
assert layer[0].conv1.out_channels == 16
assert layer[0].downsample is not None and len(layer[0].downsample) == 2
assert isinstance(layer[0].downsample[0], nn.Conv2d)
assert layer[0].downsample[0].stride == (1, 1)
for i in range(1, 3):
assert layer[i].in_channels == 64
assert layer[i].out_channels == 64
assert layer[i].conv1.out_channels == 16
assert layer[i].stride == 1
assert layer[i].downsample is None
x = torch.randn(1, 32, 56, 56)
x_out = layer(x)
assert x_out.shape == (1, 64, 56, 56)
# 3 Bottleneck w/ stride 2 and downsample
layer = ResLayer(Bottleneck, 3, 32, 64, stride=2)
assert len(layer) == 3
assert layer[0].in_channels == 32
assert layer[0].out_channels == 64
assert layer[0].stride == 2
assert layer[0].conv1.out_channels == 16
assert layer[0].downsample is not None and len(layer[0].downsample) == 2
assert isinstance(layer[0].downsample[0], nn.Conv2d)
assert layer[0].downsample[0].stride == (2, 2)
for i in range(1, 3):
assert layer[i].in_channels == 64
assert layer[i].out_channels == 64
assert layer[i].conv1.out_channels == 16
assert layer[i].stride == 1
assert layer[i].downsample is None
x = torch.randn(1, 32, 56, 56)
x_out = layer(x)
assert x_out.shape == (1, 64, 28, 28)
# 3 Bottleneck w/ stride 2 and downsample with avg pool
layer = ResLayer(Bottleneck, 3, 32, 64, stride=2, avg_down=True)
assert len(layer) == 3
assert layer[0].in_channels == 32
assert layer[0].out_channels == 64
assert layer[0].stride == 2
assert layer[0].conv1.out_channels == 16
assert layer[0].downsample is not None and len(layer[0].downsample) == 3
assert isinstance(layer[0].downsample[0], nn.AvgPool2d)
assert layer[0].downsample[0].stride == 2
for i in range(1, 3):
assert layer[i].in_channels == 64
assert layer[i].out_channels == 64
assert layer[i].conv1.out_channels == 16
assert layer[i].stride == 1
assert layer[i].downsample is None
x = torch.randn(1, 32, 56, 56)
x_out = layer(x)
assert x_out.shape == (1, 64, 28, 28)
# 3 Bottleneck with custom expansion
layer = ResLayer(Bottleneck, 3, 32, 32, expansion=2)
assert len(layer) == 3
for i in range(3):
assert layer[i].in_channels == 32
assert layer[i].out_channels == 32
assert layer[i].stride == 1
assert layer[i].conv1.out_channels == 16
assert layer[i].downsample is None
x = torch.randn(1, 32, 56, 56)
x_out = layer(x)
assert x_out.shape == (1, 32, 56, 56)