def __init__(self, inplanes, planes, stride=1):
super(Bottleneck, self).__init__()
self.inplanes = inplanes
self.planes = planes
self.conv1 = Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = FrozenBatchNorm2d(planes)
self.conv2 = Conv2d(planes,
planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
self.bn2 = FrozenBatchNorm2d(planes)
self.conv3 = Conv2d(planes,
planes * self.expansion,
kernel_size=1,
bias=False)
self.bn3 = FrozenBatchNorm2d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
if self.inplanes != self.planes * self.expansion:
self.downsample = nn.Sequential(
Conv2d(self.inplanes,
self.planes * self.expansion,
kernel_size=1,
stride=stride,
bias=False),
FrozenBatchNorm2d(self.planes * self.expansion),
)
def __init__(self, C):
super(SEModule, self).__init__()
mid = max(C // self.reduction, 8)
conv1 = Conv2d(C, mid, 1, 1, 0)
conv2 = Conv2d(mid, C, 1, 1, 0)
self.op = nn.Sequential(nn.AdaptiveAvgPool2d(1), conv1,
nn.ReLU(inplace=True), conv2, nn.Sigmoid())
def _make_fuse_layers(self):
if self.num_branches == 1:
return None
num_branches = self.num_branches
num_inchannels = self.num_inchannels
fuse_layers = []
for i in range(num_branches if self.multi_scale_output else 1):
fuse_layer = []
for j in range(num_branches):
if j > i:
fuse_layer.append(
nn.Sequential(
Conv2d(num_inchannels[j],
num_inchannels[i],
1,
1,
0,
bias=False),
FrozenBatchNorm2d(num_inchannels[i]),
nn.Upsample(scale_factor=2**(j - i),
mode='nearest')))
elif j == i:
fuse_layer.append(None)
else:
conv3x3s = []
for k in range(i - j):
if k == i - j - 1:
num_outchannels_conv3x3 = num_inchannels[i]
conv3x3s.append(
nn.Sequential(
Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3,
2,
1,
bias=False),
FrozenBatchNorm2d(
num_outchannels_conv3x3)))
else:
num_outchannels_conv3x3 = num_inchannels[j]
conv3x3s.append(
nn.Sequential(
Conv2d(num_inchannels[j],
num_outchannels_conv3x3,
3,
2,
1,
bias=False),
FrozenBatchNorm2d(num_outchannels_conv3x3),
nn.ReLU(True)))
fuse_layer.append(nn.Sequential(*conv3x3s))
fuse_layers.append(nn.ModuleList(fuse_layer))
return nn.ModuleList(fuse_layers)
def __init__(self, C_in, C_out, stride):
assert stride in [1, 2]
ops = [
Conv2d(C_in, C_in, 3, stride, 1, bias=False),
BatchNorm2d(C_in),
nn.ReLU(inplace=True),
Conv2d(C_in, C_out, 3, 1, 1, bias=False),
BatchNorm2d(C_out),
]
super(CascadeConv3x3, self).__init__(*ops)
self.res_connect = (stride == 1) and (C_in == C_out)
def __init__(self, cfg, in_channels):
super(KeypointRCNNFeatureExtractor, self).__init__()
resolution = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_RESOLUTION
scales = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_SCALES
sampling_ratio = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_SAMPLING_RATIO
pooler = Pooler(
output_size=(resolution, resolution),
scales=scales,
sampling_ratio=sampling_ratio,
)
self.pooler = pooler
input_features = in_channels
layers = cfg.MODEL.ROI_KEYPOINT_HEAD.CONV_LAYERS
next_feature = input_features
self.blocks = []
for layer_idx, layer_features in enumerate(layers, 1):
layer_name = "conv_fcn{}".format(layer_idx)
module = Conv2d(next_feature,
layer_features,
3,
stride=1,
padding=1)
nn.init.kaiming_normal_(module.weight,
mode="fan_out",
nonlinearity="relu")
nn.init.constant_(module.bias, 0)
self.add_module(layer_name, module)
next_feature = layer_features
self.blocks.append(layer_name)
self.out_channels = layer_features
def __init__(self, C_in, C_out, expansion, stride):
assert stride in [1, 2]
self.res_connect = (stride == 1) and (C_in == C_out)
C_mid = _get_divisible_by(C_in * expansion, 8, 8)
ops = [
# pw
Conv2d(C_in, C_mid, 1, 1, 0, bias=False),
BatchNorm2d(C_mid),
nn.ReLU(inplace=True),
# shift
Shift(C_mid, 5, stride, 2),
# pw-linear
Conv2d(C_mid, C_out, 1, 1, 0, bias=False),
BatchNorm2d(C_out),
]
super(ShiftBlock5x5, self).__init__(*ops)
def _make_transition_layer(self, num_channels_pre_layer,
num_channels_cur_layer):
num_branches_cur = len(num_channels_cur_layer)
num_branches_pre = len(num_channels_pre_layer)
transition_layers = []
for i in range(num_branches_cur):
if i < num_branches_pre:
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
transition_layers.append(
nn.Sequential(
Conv2d(num_channels_pre_layer[i],
num_channels_cur_layer[i],
3,
1,
1,
bias=False),
FrozenBatchNorm2d(num_channels_cur_layer[i]),
nn.ReLU(inplace=True)))
else:
transition_layers.append(None)
else:
conv3x3s = []
for j in range(i + 1 - num_branches_pre):
inchannels = num_channels_pre_layer[-1]
outchannels = num_channels_cur_layer[i] \
if j == i - num_branches_pre else inchannels
conv3x3s.append(
nn.Sequential(
Conv2d(inchannels,
outchannels,
3,
2,
1,
bias=False), FrozenBatchNorm2d(outchannels),
nn.ReLU(inplace=True)))
transition_layers.append(nn.Sequential(*conv3x3s))
return nn.ModuleList(transition_layers)
def __init__(self,
input_depth,
output_depth,
kernel,
stride,
pad,
no_bias,
use_relu,
bn_type,
group=1,
*args,
**kwargs):
super(ConvBNRelu, self).__init__()
assert use_relu in ["relu", None]
if isinstance(bn_type, (list, tuple)):
assert len(bn_type) == 2
assert bn_type[0] == "gn"
gn_group = bn_type[1]
bn_type = bn_type[0]
assert bn_type in ["bn", "af", "gn", None]
assert stride in [1, 2, 4]
op = Conv2d(input_depth,
output_depth,
kernel_size=kernel,
stride=stride,
padding=pad,
bias=not no_bias,
groups=group,
*args,
**kwargs)
nn.init.kaiming_normal_(op.weight, mode="fan_out", nonlinearity="relu")
if op.bias is not None:
nn.init.constant_(op.bias, 0.0)
self.add_module("conv", op)
if bn_type == "bn":
bn_op = BatchNorm2d(output_depth)
elif bn_type == "gn":
bn_op = nn.GroupNorm(num_groups=gn_group,
num_channels=output_depth)
elif bn_type == "af":
bn_op = FrozenBatchNorm2d(output_depth)
if bn_type is not None:
self.add_module("bn", bn_op)
if use_relu == "relu":
self.add_module("relu", nn.ReLU(inplace=True))
def __init__(self, cfg, norm_func):
super(BaseStem, self).__init__()
out_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS
self.conv1 = Conv2d(3,
out_channels,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = norm_func(out_channels)
for l in [
self.conv1,
]:
nn.init.kaiming_uniform_(l.weight, a=1)
def __init__(
self,
in_channels,
num_levels,
refine_level=2,
refine_type=None,
):
super(BFP, self).__init__()
assert refine_type in [None, 'conv', 'non_local']
assert 0 <= refine_level < num_levels
self.in_channels = in_channels
self.num_levels = num_levels
self.refine_level = refine_level
self.refine_type = refine_type
if self.refine_type == 'conv':
self.refine = nn.Sequential(
Conv2d(self.in_channels,
self.in_channels,
kernel_size=3,
stride=1,
padding=1),
BatchNorm2d(self.in_channels),
nn.ReLU(inplace=True),
)
for m in self.refine:
if isinstance(m, Conv2d):
nn.init.kaiming_normal_(m.weight)
elif isinstance(m, BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif self.refine_type == 'non_local':
self.refine = NonLocal2D(
self.in_channels,
reduction=1,
use_scale=False,
)
def __init__(
self,
in_channels,
bottleneck_channels,
out_channels,
num_groups,
stride_in_1x1,
stride,
dilation,
norm_func,
dcn_config,
gcb_config,
cbam_config,
):
super(Bottleneck, self).__init__()
self.downsample = None
if in_channels != out_channels:
down_stride = stride if dilation == 1 else 1
self.downsample = nn.Sequential(
Conv2d(in_channels,
out_channels,
kernel_size=1,
stride=down_stride,
bias=False),
norm_func(out_channels),
)
for modules in [
self.downsample,
]:
for l in modules.modules():
if isinstance(l, Conv2d):
nn.init.kaiming_uniform_(l.weight, a=1)
if dilation > 1:
stride = 1 # reset to be 1
# The original MSRA ResNet models have stride in the first 1x1 conv
# The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
# stride in the 3x3 conv
stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
self.conv1 = Conv2d(
in_channels,
bottleneck_channels,
kernel_size=1,
stride=stride_1x1,
bias=False,
)
self.bn1 = norm_func(bottleneck_channels)
# TODO: specify init for the above
with_dcn = dcn_config.get("stage_with_dcn", False)
if with_dcn:
deformable_groups = dcn_config.get("deformable_groups", 1)
with_modulated_dcn = dcn_config.get("with_modulated_dcn", False)
self.conv2 = DFConv2d(bottleneck_channels,
bottleneck_channels,
with_modulated_dcn=with_modulated_dcn,
kernel_size=3,
stride=stride_3x3,
groups=num_groups,
dilation=dilation,
deformable_groups=deformable_groups,
bias=False)
else:
self.conv2 = Conv2d(bottleneck_channels,
bottleneck_channels,
kernel_size=3,
stride=stride_3x3,
padding=dilation,
bias=False,
groups=num_groups,
dilation=dilation)
nn.init.kaiming_uniform_(self.conv2.weight, a=1)
self.bn2 = norm_func(bottleneck_channels)
self.conv3 = Conv2d(bottleneck_channels,
out_channels,
kernel_size=1,
bias=False)
self.bn3 = norm_func(out_channels)
for l in [
self.conv1,
self.conv3,
]:
nn.init.kaiming_uniform_(l.weight, a=1)
## add global context block
## this module is initialized when object created
## no need to re-initialize the weight
self.with_gcb = gcb_config.get("stage_with_gcb", False)
if self.with_gcb:
self.context_block = GlobalContextBlock(
inplanes=out_channels,
ratio=gcb_config["ratio"],
pooling_type=gcb_config["pooling_type"],
fusion_types=gcb_config["fusion_types"],
)
## add cbam block
self.with_cbam = cbam_config.get("stage_with_cbam", False)
if self.with_cbam:
self.cbam_block = CBAM(
gate_channels=out_channels,
reduction_ratio=cbam_config["reduction_ratio"],
pool_types=cbam_config["pool_types"],
no_spatial=cbam_config['no_spatial'])