def get_extras(source_out_channels,
use_depthwise=False,
layers=(),
activation={
'name': 'ReLU',
'args': {
'inplace': True
}
},
initializer={'name': 'xavier_normal_'},
batch_norm={}):
extras = nn.ModuleList()
in_channels = source_out_channels[-1]
extra_layers = layers
for type_, out_channels in extra_layers:
layers = []
if type_ == 'm':
out_channels = in_channels
layers.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
elif type_ == 's':
layers.append(
conv.Conv2dBn(in_channels,
out_channels // 2,
kernel_size=1,
bias=False,
activation_params=activation,
use_bn=True,
batch_norm_params=batch_norm))
in_channels = out_channels // 2
if use_depthwise:
layers.append(
conv.DepthwiseConv2dBn(in_channels,
out_channels,
kernel_size=3,
stride=2,
padding=1,
bias=False,
activation_params=activation,
use_bn=True,
batch_norm_params=batch_norm))
else:
layers.append(
conv.Conv2dBn(in_channels,
out_channels,
kernel_size=3,
stride=2,
padding=1,
bias=False,
activation_params=activation,
use_bn=True,
batch_norm_params=batch_norm))
elif type_ == '':
layers.append(
conv.Conv2dBn(in_channels,
out_channels // 2,
kernel_size=1,
bias=False,
activation_params=activation,
use_bn=True,
batch_norm_params=batch_norm))
in_channels = out_channels // 2
if use_depthwise:
layers.append(
conv.DepthwiseConv2dBn(in_channels,
out_channels,
kernel_size=3,
bias=False,
activation_params=activation,
use_bn=True,
batch_norm_params=batch_norm))
else:
layers.append(
conv.Conv2dBn(in_channels,
out_channels,
kernel_size=3,
bias=False,
activation_params=activation,
use_bn=True,
batch_norm_params=batch_norm))
else:
raise ValueError(f'Unknown layer type: {type_}')
source_out_channels.append(out_channels)
extras.append(nn.Sequential(*layers))
in_channels = out_channels
initializer_ = functools.partial(getattr(nn.init, initializer['name']),
**initializer.get('args', {}))
def _init_extras(layer):
if isinstance(layer, nn.Conv2d):
initializer_(layer.weight)
layer.bias is not None and nn.init.zeros_(layer.bias)
extras.apply(_init_extras)
return extras
def __init__(self,
base,
out_layers,
pyramid_layers,
pyramid_channels,
interpolation_mode='nearest',
activation={
'name': 'ReLU',
'args': {
'inplace': True
}
},
initializer={'name': 'xavier_normal_'},
**kwargs):
super(DepthwiseFeaturePyramid, self).__init__(base, out_layers,
**kwargs)
self.pyramid_layers = pyramid_layers
self.pyramid_channels = pyramid_channels
self.interpolation_mode = interpolation_mode
self.pyramid_lateral = nn.ModuleList()
self.downsample = nn.ModuleList()
self.up_conv = nn.ModuleList()
self.num_outputs = pyramid_layers
base_out_channels = super(DepthwiseFeaturePyramid,
self).get_out_channels()
for in_channels in base_out_channels:
self.pyramid_lateral.append(
nn.Conv2d(in_channels, pyramid_channels, kernel_size=1))
for _ in range(pyramid_layers - len(out_layers)):
paths = nn.ModuleList()
paths.append(
nn.Sequential(
nn.MaxPool2d(kernel_size=2),
conv.Conv2dBn(pyramid_channels,
pyramid_channels // 2,
kernel_size=1,
activation_params=activation)))
paths.append(
conv.DepthwiseConv2dBn(pyramid_channels,
pyramid_channels // 2,
kernel_size=3,
stride=2,
padding=1,
activation_params=activation))
self.downsample.append(paths)
for _ in range(pyramid_layers - 1):
self.up_conv.append(
conv.Conv2dBn(pyramid_channels,
pyramid_channels,
kernel_size=3,
padding=1,
groups=pyramid_channels,
activation_params=activation))
self.pyramid_lateral.apply(self.init_layer)
self.downsample.apply(self.init_layer)
self.up_conv.apply(self.init_layer)
def get_predictor(source_out_channels,
num_boxes,
num_classes,
use_depthwise=False,
num_layers=0,
num_channels=256,
kernel_size=3,
batch_norm={},
activation={
'name': 'ReLU',
'args': {
'inplace': True
}
},
initializer={
'name': 'normal_',
'args': {
'mean': 0,
'std': 0.01
}
},
class_head_bias_init=0):
if num_layers > 0:
assert len(set(source_out_channels)) == 1
predictor = nn.ModuleDict()
norms = nn.ModuleDict()
for head in ['class', 'loc']:
in_channels = source_out_channels[0]
layers = nn.ModuleList()
norms[head] = nn.ModuleList()
for _ in range(num_layers):
if use_depthwise:
layers.append(
conv.DepthwiseConv2dBn(in_channels,
num_channels,
kernel_size=kernel_size,
padding=1,
bias=True,
activation_params=None,
use_bn=False))
else:
layers.append(
conv.Conv2dBn(in_channels,
num_channels,
kernel_size=kernel_size,
padding=1,
bias=True,
activation_params=None,
use_bn=False))
layer_norms = nn.ModuleList()
for _ in source_out_channels:
layer_norms.append(nn.BatchNorm2d(num_channels, **batch_norm))
norms[head].append(layer_norms)
in_channels = num_channels
predictor[head] = layers
activation_ = getattr(nn, activation['name'])(**activation.get('args', {}))
if num_layers > 0:
out_channels = [num_channels] * len(source_out_channels)
else:
out_channels = source_out_channels
initializer_ = functools.partial(getattr(nn.init, initializer['name']),
**initializer.get('args', {}))
def _init_predictor(layer):
if isinstance(layer, nn.Conv2d):
initializer_(layer.weight)
nn.init.zeros_(layer.bias)
predictor.apply(_init_predictor)
def _init_class_head(layer):
initializer_(layer.weight)
nn.init.constant_(layer.bias, class_head_bias_init)
heads = nn.ModuleList()
for in_channels, num_boxes in zip(out_channels, num_boxes):
class_head = nn.Conv2d(in_channels,
num_boxes * num_classes,
kernel_size=3,
padding=1,
bias=True)
loc_head = nn.Conv2d(in_channels,
num_boxes * 4,
kernel_size=3,
padding=1,
bias=True)
class_head.apply(_init_class_head)
loc_head.apply(_init_predictor)
heads.append(nn.ModuleDict({'class': class_head, 'loc': loc_head}))
return (predictor, activation_, norms), heads
def __init__(self,
source_out_channels,
num_boxes,
num_classes,
use_depthwise,
num_layers=0,
num_channels=256,
kernel_size=3,
batch_norm={},
activation={
'name': 'ReLU',
'args': {
'inplace': True
}
},
initializer={
'name': 'normal_',
'args': {
'mean': 0,
'std': 0.01
}
}):
super(SharedConvPredictor, self).__init__()
if num_layers > 0:
assert len(set(source_out_channels)) == 1
self.convs = nn.ModuleDict()
self.norms = nn.ModuleDict()
for head in ['score', 'loc']:
in_channels = source_out_channels[0]
layers = nn.ModuleList()
self.norms[head] = nn.ModuleList()
for _ in range(num_layers):
if use_depthwise:
layers.append(
conv.DepthwiseConv2dBn(in_channels,
num_channels,
kernel_size=kernel_size,
padding=1,
bias=True,
activation_params=None,
use_bn=False))
else:
layers.append(
conv.Conv2dBn(in_channels,
num_channels,
kernel_size=kernel_size,
padding=1,
bias=True,
activation_params=None,
use_bn=False))
layer_norms = nn.ModuleList()
for _ in source_out_channels:
layer_norms.append(
nn.BatchNorm2d(num_channels, **batch_norm))
self.norms[head].append(layer_norms)
in_channels = num_channels
self.convs[head] = layers
self.activation = getattr(
nn, activation['name'])(**activation.get('args', {}))
self.out_channels = [num_channels] * len(source_out_channels)
initializer_ = functools.partial(getattr(nn.init, initializer['name']),
**initializer.get('args', {}))
def _init_predictor(layer):
if isinstance(layer, nn.Conv2d):
initializer_(layer.weight)
nn.init.zeros_(layer.bias)
self.convs.apply(_init_predictor)