in scripts/models/wider_resnet.py [0:0]
def __init__(self, structure, norm_act=ABN, classes=0):
"""Wider ResNet with pre-activation (identity mapping) blocks
Parameters
----------
structure : list of int
Number of residual blocks in each of the six modules of the network.
norm_act : callable
Function to create normalization / activation Module.
classes : int
If not `0` also include global average pooling and a fully-connected layer with `classes` outputs at the end
of the network.
"""
super(WiderResNet, self).__init__()
self.structure = structure
if len(structure) != 6:
raise ValueError("Expected a structure with six values")
# Initial layers
self.mod1 = nn.Sequential(
OrderedDict(
[("conv1", nn.Conv2d(3, 64, 3, stride=1, padding=1, bias=False))]
)
)
# Groups of residual blocks
in_channels = 64
channels = [
(128, 128),
(256, 256),
(512, 512),
(512, 1024),
(512, 1024, 2048),
(1024, 2048, 4096),
]
for mod_id, num in enumerate(structure):
# Create blocks for module
blocks = []
for block_id in range(num):
blocks.append(
(
"block%d" % (block_id + 1),
IdentityResidualBlock(
in_channels, channels[mod_id], norm_act=norm_act
),
)
)
# Update channels and p_keep
in_channels = channels[mod_id][-1]
# Create module
if mod_id <= 4:
self.add_module(
"pool%d" % (mod_id + 2), nn.MaxPool2d(3, stride=2, padding=1)
)
self.add_module("mod%d" % (mod_id + 2), nn.Sequential(OrderedDict(blocks)))
# Pooling and predictor
self.bn_out = norm_act(in_channels)
if classes != 0:
self.classifier = nn.Sequential(
OrderedDict(
[
("avg_pool", GlobalAvgPool2d()),
("fc", nn.Linear(in_channels, classes)),
]
)
)