def _make_stem_layer(self):
"""Make stem layer."""
stem_channel = self.base_channel // 2
norm_cfg = self.norm_cfg.copy()
if self.norm_cfg.get('type') == 'GN':
num_groups = norm_cfg.get('num_groups')
norm_cfg['num_groups'] = int(num_groups / 2)
norm_creator = LayerCreator(**norm_cfg)
conv_creator = LayerCreator(**self.conv_cfg)
self.conv1 = conv_creator.create_layer(3,
stem_channel,
kernel_size=3,
stride=2,
padding=1,
bias=False)
self.norm1_name = norm_creator.get_name(magic_number=1)
norm1 = norm_creator.create_layer(num_features=stem_channel)
self.add_module(self.norm1_name, norm1)
self.conv2 = conv_creator.create_layer(stem_channel,
self.base_channel,
kernel_size=3,
stride=2,
padding=1,
bias=False)
self.norm2_name = norm_creator.get_name(magic_number=2)
norm2 = norm_creator.create_layer(num_features=self.base_channel)
self.add_module(self.norm2_name, norm2)
self.relu = nn.ReLU(inplace=True)
def __init__(self,
inplanes,
planes,
stride=1,
dilation=1,
downsample=None,
style='pytorch',
conv_cfg=None,
norm_cfg=None, **kwargs):
super(BasicBlock, self).__init__()
norm_creator = LayerCreator(**norm_cfg)
self.norm1_name = norm_creator.get_name(magic_number=1)
norm1 = norm_creator.create_layer(num_features=planes)
self.add_module(self.norm1_name, norm1)
self.norm2_name = norm_creator.get_name(magic_number=2)
norm2 = norm_creator.create_layer(num_features=planes)
conv_creator = LayerCreator(**conv_cfg)
self.conv1 = conv_creator.create_layer(inplanes, planes, 3, stride=stride, padding=dilation,
dilation=dilation, bias=False)
self.conv2 = conv_creator.create_layer(planes, planes, 3, padding=1, bias=False)
self.add_module(self.norm2_name, norm2)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
self.dilation = dilation
def make_res_layer(block,
inplanes,
planes,
arch,
stride=1,
dilation=1,
style='pytorch',
conv_cfg=None,
norm_cfg=None, ):
"""Make res layer.
:param block: block function
:type block: nn.Module
:param inplanes: input feature map channel num
:type inplanes: int
:param planes: output feature map channel num
:type planes: int
:param arch: model arch
:type arch: list
:param stride: stride
:type stride: int
:param dilation: dilation
:type dilation: int
:param style: style
:type style: str
:param conv_cfg: conv config
:type conv_cfg: dict
:param norm_cfg: norm config
:type norm_cfg: dict
:return: res layer
:rtype: nn.Module
"""
conv_creator = LayerCreator(**conv_cfg)
norm_creator = LayerCreator(**norm_cfg)
layers = []
for i, layer_type in enumerate(arch):
downsample = None
stride = stride if i == 0 else 1
if layer_type == 2:
planes *= 2
if stride != 1 or inplanes != planes * block.expansion:
downsample = nn.Sequential(
conv_creator.create_layer(in_channels=inplanes, out_channels=planes * block.expansion, kernel_size=1,
stride=stride, bias=False),
norm_creator.create_layer(num_features=planes * block.expansion))
layers.append(block(inplanes=inplanes, planes=planes, stride=stride, dilation=dilation, downsample=downsample,
style=style, conv_cfg=conv_cfg, norm_cfg=norm_cfg))
inplanes = planes * block.expansion
return nn.Sequential(*layers)
def __init__(self,
inplanes,
planes,
stride=1,
dilation=1,
downsample=None,
style='pytorch',
conv_cfg=None,
norm_cfg=None,
**kwargs):
super(Bottleneck, self).__init__()
self.inplanes = inplanes
self.planes = planes
self.stride = stride
self.dilation = dilation
self.style = style
self.conv_cfg = conv_cfg
self.norm_cfg = norm_cfg
if self.style == 'pytorch':
self.conv1_stride = 1
self.conv2_stride = stride
else:
self.conv1_stride = stride
self.conv2_stride = 1
norm_creator = LayerCreator(**norm_cfg)
self.norm1_name = norm_creator.get_name(magic_number=1)
norm1 = norm_creator.create_layer(num_features=planes)
self.norm2_name = norm_creator.get_name(magic_number=2)
norm2 = norm_creator.create_layer(num_features=planes)
self.norm3_name = norm_creator.get_name(magic_number=3)
norm3 = norm_creator.create_layer(num_features=planes * self.expansion)
conv_creator = LayerCreator(**conv_cfg)
self.conv1 = conv_creator.create_layer(inplanes,
planes,
kernel_size=1,
stride=self.conv1_stride,
bias=False)
self.add_module(self.norm1_name, norm1)
self.conv2 = conv_creator.create_layer(planes,
planes,
kernel_size=3,
stride=self.conv2_stride,
padding=dilation,
dilation=dilation,
bias=False)
self.add_module(self.norm2_name, norm2)
self.conv3 = conv_creator.create_layer(planes,
planes * self.expansion,
kernel_size=1,
bias=False)
self.add_module(self.norm3_name, norm3)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample