def __init__(self, builder: ConvBuilder, in_planes, stage_deps, stride=1):
super(ResNetBottleneckStage, self).__init__()
print('building stage: in {}, deps {}'.format(in_planes, stage_deps))
assert (len(stage_deps) - 1) % 3 == 0
self.num_blocks = (len(stage_deps) - 1) // 3
stage_out_channels = stage_deps[3]
for i in range(2, self.num_blocks):
assert stage_deps[3 * i] == stage_out_channels
self.relu = builder.ReLU()
self.projection = builder.Conv2dBN(in_channels=in_planes,
out_channels=stage_deps[0],
kernel_size=1,
stride=stride)
self.align_opr = builder.ResNetAlignOpr(channels=stage_deps[0])
for i in range(self.num_blocks):
in_c = in_planes if i == 0 else stage_out_channels
block_stride = stride if i == 0 else 1
self.__setattr__(
'block{}'.format(i),
BottleneckBranch(builder=builder,
in_channels=in_c,
deps=stage_deps[1 + i * 3:4 + i * 3],
stride=block_stride))
def __init__(self,
builder: ConvBuilder,
in_planes,
stage_deps,
stride=1,
is_first=False):
super(ResNetBasicStage, self).__init__()
print('building stage: in {}, deps {}'.format(in_planes, stage_deps))
self.num_blocks = len(stage_deps) // 2
stage_out_channels = stage_deps[0]
for i in range(0, self.num_blocks):
assert stage_deps[i * 2 + 2] == stage_out_channels
if is_first:
self.conv1 = builder.Conv2dBN(in_channels=in_planes,
out_channels=stage_out_channels,
kernel_size=3,
stride=1,
padding=1)
# self.projection = builder.ResIdentity(num_channels=stage_out_channels)
else:
self.projection = builder.Conv2dBN(in_channels=in_planes,
out_channels=stage_out_channels,
kernel_size=1,
stride=stride)
self.relu = builder.ReLU()
self.align_opr = builder.ResNetAlignOpr(channels=stage_out_channels)
for i in range(self.num_blocks):
if i == 0 and is_first:
in_c = stage_deps[0]
elif i == 0:
in_c = in_planes
else:
in_c = stage_out_channels
block_stride = stride if i == 0 else 1
self.__setattr__(
'block{}'.format(i),
BasicBranch(builder=builder,
in_channels=in_c,
deps=stage_deps[1 + i * 2:3 + i * 2],
stride=block_stride))
