def make_conv3x3(
in_channels,
out_channels,
dilation=1,
stride=1,
use_gn=False,
use_relu=False,
kaiming_init=True
):
conv = Conv2d(
in_channels,
out_channels,
kernel_size=3,
stride=stride,
padding=dilation,
dilation=dilation,
bias=False if use_gn else True
)
if kaiming_init:
init.kaiming_normal_(
conv.weight, mode="fan_out", nonlinearity="relu"
)
else:
init.gauss_(conv.weight, std=0.01)
if not use_gn:
init.constant_(conv.bias, 0)
module = [conv,]
if use_gn:
module.append(group_norm(out_channels))
if use_relu:
module.append(nn.ReLU())
if len(module) > 1:
return nn.Sequential(*module)
return conv
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)
init.kaiming_normal_(module.weight,
mode="fan_out",
nonlinearity="relu")
init.constant_(module.bias, 0)
setattr(self, layer_name, module)
next_feature = layer_features
self.blocks.append(layer_name)
self.out_channels = layer_features
def __init__(self, num_features):
super().__init__()
self.num_features = num_features
self.weight = jt.zeros((num_features))
self.bias = jt.zeros((num_features))
init.uniform_(self.weight, 0, 1)
init.constant_(self.bias, 0.0)
def __init__(self, cfg, in_channels, num_anchors):
"""
Arguments:
cfg : config
in_channels (int): number of channels of the input feature
num_anchors (int): number of anchors to be predicted
"""
super(RPNHead, self).__init__()
self.conv = nn.Conv(in_channels,
in_channels,
kernel_size=3,
stride=1,
padding=1)
self.cls_logits = nn.Conv(in_channels,
num_anchors,
kernel_size=1,
stride=1)
self.bbox_pred = nn.Conv(in_channels,
num_anchors * 4,
kernel_size=1,
stride=1)
for l in [self.conv, self.cls_logits, self.bbox_pred]:
init.gauss_(l.weight, std=0.01)
init.constant_(l.bias, 0)
def __init__(self, cfg, in_channels):
"""
Arguments:
in_channels (int): number of channels of the input feature
num_anchors (int): number of anchors to be predicted
"""
super(RetinaNetHead, self).__init__()
# TODO: Implement the sigmoid version first.
num_classes = cfg.MODEL.RETINANET.NUM_CLASSES - 1
num_anchors = len(cfg.MODEL.RETINANET.ASPECT_RATIOS) \
* cfg.MODEL.RETINANET.SCALES_PER_OCTAVE
cls_tower = []
bbox_tower = []
for i in range(cfg.MODEL.RETINANET.NUM_CONVS):
cls_tower.append(
nn.Conv(
in_channels,
in_channels,
kernel_size=3,
stride=1,
padding=1
)
)
cls_tower.append(nn.ReLU())
bbox_tower.append(
nn.Conv(
in_channels,
in_channels,
kernel_size=3,
stride=1,
padding=1
)
)
bbox_tower.append(nn.ReLU())
self.cls_tower = nn.Sequential(*cls_tower)
self.bbox_tower = nn.Sequential(*bbox_tower)
self.cls_logits = nn.Conv(
in_channels, num_anchors * num_classes, kernel_size=3, stride=1,
padding=1
)
self.bbox_pred = nn.Conv(
in_channels, num_anchors * 4, kernel_size=3, stride=1,
padding=1
)
# Initialization
for modules in [self.cls_tower, self.bbox_tower, self.cls_logits,
self.bbox_pred]:
for l in modules.modules():
if isinstance(l, nn.Conv):
init.gauss_(l.weight, std=0.01)
init.constant_(l.bias, 0)
# retinanet_bias_init
prior_prob = cfg.MODEL.RETINANET.PRIOR_PROB
bias_value = -math.log((1 - prior_prob) / prior_prob)
init.constant_(self.cls_logits.bias, bias_value)
def __init__(self, block, layers, baseWidth=26, scale=4, num_classes=1000):
self.inplanes = 64
super(Res2Net, self).__init__()
self.baseWidth = baseWidth
self.scale = scale
self.conv1 = nn.Sequential(
nn.Conv(3, 32, 3, stride=2, padding=1, bias=False),
nn.BatchNorm(32), nn.ReLU(),
nn.Conv(32, 32, 3, stride=1, padding=1, bias=False),
nn.BatchNorm(32), nn.ReLU(),
nn.Conv(32, 64, 3, stride=1, padding=1, bias=False))
self.bn1 = nn.BatchNorm(64)
self.relu = nn.ReLU()
self.maxpool = nn.Pool(3, stride=2, padding=1, op='maximum')
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Linear((512 * block.expansion), num_classes)
for m in self.modules():
if isinstance(m, nn.Conv):
nn.init.kaiming_normal_(m.weight, mode='fan_out')
elif isinstance(m, nn.BatchNorm):
init.constant_(m.weight, value=1)
init.constant_(m.bias, value=0)
def weights_init_normal(m):
classname = m.__class__.__name__
if (classname.find('Linear') != (- 1)):
init.gauss_(m.weight, mean=0.0, std=0.02)
elif (classname.find('BatchNorm') != (- 1)):
init.gauss_(m.weight, mean=1.0, std=0.02)
init.constant_(m.bias, value=0.0)
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
init.gauss_(m.weight, 0.0, 0.02)
init.gauss_(m.bias, 0.0, 0.02)
elif classname.find("BatchNorm") != -1:
init.gauss_(m.weight, 1.0, 0.02)
init.constant_(m.bias, 0.0)
def __init__(self, cfg):
super(MaskIoUPredictor, self).__init__()
num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
self.maskiou = nn.Linear(1024, num_classes)
init.gauss_(self.maskiou.weight, mean=0, std=0.01)
init.constant_(self.maskiou.bias, 0)
def __init__(self, in_channels, out_channels):
super(LastLevelP6P7, self).__init__()
self.p6 = nn.Conv(in_channels, out_channels, 3, 2, 1)
self.p7 = nn.Conv(out_channels, out_channels, 3, 2, 1)
for module in [self.p6, self.p7]:
init.kaiming_uniform_(module.weight, a=1)
init.constant_(module.bias, 0)
self.use_P5 = in_channels == out_channels
def test_constant(self):
a = init.constant(2, "float32")
np.testing.assert_allclose(a.data, [0, 0])
a = init.constant((2, 3), value=1.)
np.testing.assert_allclose(a.data, [[1, 1, 1], [1, 1, 1]])
linear = nn.Linear(2, 2)
init.constant_(linear.weight)
np.testing.assert_allclose(linear.weight.data, [[0, 0], [0, 0]])
def __init__(self, n_classes):
super(SSD300, self).__init__()
self.n_classes = n_classes
self.base = VGGBase()
self.aux_convs = AuxiliaryConvolutions()
self.pred_convs = PredictionConvolutions(n_classes)
self.rescale_factors = jt.make_var([1, 512, 1, 1], init=jt.zeros)
init.constant_(self.rescale_factors, 20)
self.priors_cxcy = self.create_prior_boxes()
def __init__(self, cfg):
super(MaskIoUGAPPredictor, self).__init__()
num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
in_channel = cfg.MODEL.ROI_MASK_HEAD.CONV_LAYERS[0]
self.maskiou = nn.Linear(in_channel, num_classes)
init.gauss_(self.maskiou.weight, mean=0, std=0.01)
init.constant_(self.maskiou.bias, 0)
def weights_init_normal(m):
classname = m.__class__.__name__
if (classname.find('Conv') != (-1)):
init.gauss_(m.weight, mean=0.0, std=0.02)
if (hasattr(m, 'bias') and (m.bias is not None)):
init.constant_(m.bias, value=0.0)
elif (classname.find('BatchNorm') != (-1)):
init.gauss_(m.weight, mean=1.0, std=0.02)
init.constant_(m.bias, value=0.0)
def make_fc(dim_in, hidden_dim, use_gn=False):
if use_gn:
fc = nn.Linear(dim_in, hidden_dim, bias=False)
init.kaiming_uniform_(fc.weight, a=1)
return nn.Sequential(fc, group_norm(hidden_dim))
fc = nn.Linear(dim_in, hidden_dim)
init.kaiming_uniform_(fc.weight, a=1)
init.constant_(fc.bias, 0)
return fc
def init_weights(m, mode='MSRAFill'):
if isinstance(m, (nn.Conv, nn.ConvTranspose)):
if mode == 'GaussianFill':
init.gauss_(m.weight, std=0.001)
elif mode == 'MSRAFill':
MSRAFill(m.weight)
else:
raise ValueError
if m.bias is not None:
init.constant_(m.bias, 0)
if isinstance(m, nn.Linear):
XavierFill(m.weight)
init.constant_(m.bias, 0)
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
init.gauss_(m.weight, 0, (2. / n)**0.5)
if m.bias is not None:
init.constant_(m.bias, 0.0)
# elif isinstance(m, BatchNorm2d):
# m.weight.data.fill_(1)
# m.bias.data.zero_()
elif isinstance(m, nn.Linear):
n = m.weight.shape[1]
init.gauss_(m.weight, 0, 0.01)
init.constant_(m.bias, 0.0)
def __init__(self, cfg, in_channels):
super(FPNPredictor, self).__init__()
num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
representation_size = in_channels
self.cls_score = nn.Linear(representation_size, num_classes)
num_bbox_reg_classes = 2 if cfg.MODEL.CLS_AGNOSTIC_BBOX_REG else num_classes
self.bbox_pred = nn.Linear(representation_size,
num_bbox_reg_classes * 4)
init.gauss_(self.cls_score.weight, std=0.01)
init.gauss_(self.bbox_pred.weight, std=0.001)
for l in [self.cls_score, self.bbox_pred]:
init.constant_(l.bias, 0)
def __init__(self,
cardinality,
depth,
nlabels,
base_width,
widen_factor=4):
""" Constructor
Args:
cardinality: number of convolution groups.
depth: number of layers.
nlabels: number of classes
base_width: base number of channels in each group.
widen_factor: factor to adjust the channel dimensionality
"""
super(CifarResNeXt, self).__init__()
self.cardinality = cardinality
self.depth = depth
self.block_depth = (self.depth - 2) // 9
self.base_width = base_width
self.widen_factor = widen_factor
self.nlabels = nlabels
self.output_size = 64
self.stages = [
64, 64 * self.widen_factor, 128 * self.widen_factor,
256 * self.widen_factor
]
self.conv_1_3x3 = nn.Conv(3, 64, 3, 1, 1, bias=False)
self.bn_1 = nn.BatchNorm(64)
self.stage_1 = self.block('stage_1', self.stages[0], self.stages[1], 1)
self.stage_2 = self.block('stage_2', self.stages[1], self.stages[2], 2)
self.stage_3 = self.block('stage_3', self.stages[2], self.stages[3], 2)
self.classifier = nn.Linear(self.stages[3], nlabels)
self.pool = nn.Pool(8, 1, op="mean")
self.relu = nn.Relu()
init.relu_invariant_gauss_(self.classifier.weight)
for param in self.parameters():
key = param.name()
if key.split('.')[-1] == 'weight':
if 'Conv' in key:
init.relu_invariant_gauss_(param, mode='fan_out')
if 'BatchNorm' in key:
init.constant_(param, value=1.0)
elif key.split('.')[-1] == 'bias':
init.constant_(param, value=0.0)
def __init__(self, cfg, in_channels):
super(MaskRCNNConv1x1Predictor, self).__init__()
num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
num_inputs = in_channels
self.mask_fcn_logits = Conv2d(num_inputs, num_classes, 1, 1, 0)
for param in self.parameters():
name = param.name()
if "bias" in name:
init.constant_(param, 0)
elif "weight" in name:
# Caffe2 implementation uses MSRAFill, which in fact
init.kaiming_normal_(param, mode="fan_out", nonlinearity="relu")
def __init__(self, cfg, in_channels):
"""
Arguments:
cfg : config
in_channels (int): number of channels of the input feature
"""
super(RPNHeadFeatureSingleConv, self).__init__()
self.conv = nn.Conv(
in_channels, in_channels, kernel_size=3, stride=1, padding=1
)
for l in [self.conv]:
init.gauss_(l.weight, std=0.01)
init.constant_(l.bias, 0)
self.out_channels = in_channels
def __init__(self, cfg, in_channels):
super(FPNXconv1fcFeatureExtractor, self).__init__()
resolution = cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION
scales = cfg.MODEL.ROI_BOX_HEAD.POOLER_SCALES
sampling_ratio = cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
pooler = Pooler(
output_size=(resolution, resolution),
scales=scales,
sampling_ratio=sampling_ratio,
)
self.pooler = pooler
use_gn = cfg.MODEL.ROI_BOX_HEAD.USE_GN
conv_head_dim = cfg.MODEL.ROI_BOX_HEAD.CONV_HEAD_DIM
num_stacked_convs = cfg.MODEL.ROI_BOX_HEAD.NUM_STACKED_CONVS
dilation = cfg.MODEL.ROI_BOX_HEAD.DILATION
xconvs = []
for ix in range(num_stacked_convs):
xconvs.append(
nn.Conv(in_channels,
conv_head_dim,
kernel_size=3,
stride=1,
padding=dilation,
dilation=dilation,
bias=False if use_gn else True))
in_channels = conv_head_dim
if use_gn:
xconvs.append(group_norm(in_channels))
xconvs.append(nn.ReLU())
self.xconvs = nn.Sequential(*xconvs)
for modules in [
self.xconvs,
]:
for l in modules.modules():
if isinstance(l, nn.Conv):
init.gauss_(l.weight, std=0.01)
if not use_gn:
init.constant_(l.bias, 0)
input_size = conv_head_dim * resolution**2
representation_size = cfg.MODEL.ROI_BOX_HEAD.MLP_HEAD_DIM
self.fc6 = make_fc(input_size, representation_size, use_gn=False)
self.out_channels = representation_size
def __init__(self, config, in_channels):
super(FastRCNNPredictor, self).__init__()
assert in_channels is not None
num_inputs = in_channels
num_classes = config.MODEL.ROI_BOX_HEAD.NUM_CLASSES
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.cls_score = nn.Linear(num_inputs, num_classes)
num_bbox_reg_classes = 2 if config.MODEL.CLS_AGNOSTIC_BBOX_REG else num_classes
self.bbox_pred = nn.Linear(num_inputs, num_bbox_reg_classes * 4)
init.gauss_(self.cls_score.weight, mean=0, std=0.01)
init.constant_(self.cls_score.bias, 0)
init.gauss_(self.bbox_pred.weight, mean=0, std=0.001)
init.constant_(self.bbox_pred.bias, 0)
def __init__(self, cfg, in_channels):
super(KeypointRCNNPredictor, self).__init__()
input_features = in_channels
num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_CLASSES
deconv_kernel = 4
self.kps_score_lowres = layers.ConvTranspose2d(
input_features,
num_keypoints,
deconv_kernel,
stride=2,
padding=deconv_kernel // 2 - 1,
)
init.kaiming_normal_(self.kps_score_lowres.weight,
mode="fan_out",
nonlinearity="relu")
init.constant_(self.kps_score_lowres.bias, 0)
self.up_scale = 2
self.out_channels = num_keypoints
def __init__(self, cfg, in_channels):
super(MaskRCNNC4Predictor, self).__init__()
num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
dim_reduced = cfg.MODEL.ROI_MASK_HEAD.CONV_LAYERS[-1]
num_inputs = in_channels
self.conv5_mask = ConvTranspose2d(num_inputs, dim_reduced, 2, 2, 0)
self.mask_fcn_logits = Conv2d(dim_reduced, num_classes, 1, 1, 0)
for param in self.parameters():
name = param.name()
if "bias" in name:
init.constant_(param, 0)
elif "weight" in name:
# Caffe2 implementation uses MSRAFill, which in fact
# corresponds to kaiming_normal_ in PyTorch
init.kaiming_normal_(param,
mode="fan_out",
nonlinearity="relu")
def reset_parameters(self):
if self.qkv_same_dim:
# Empirically observed the convergence to be much better with
# the scaled initialization
init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
else:
init.xavier_uniform_(self.k_proj.weight)
init.xavier_uniform_(self.v_proj.weight)
init.xavier_uniform_(self.q_proj.weight)
# init.xavier_uniform_(self.out_proj.weight)
if self.out_proj.bias is not None:
init.constant_(self.out_proj.bias, 0.)
if self.bias_k is not None:
init.xavier_normal_(self.bias_k)
if self.bias_v is not None:
init.xavier_normal_(self.bias_v)
def __init__(self, in_channels, n_class, roi_size, spatial_scale,
sampling_ratio):
# n_class includes the background
super(RoIHead, self).__init__()
self.classifier = nn.Sequential(
nn.Linear(in_channels * roi_size * roi_size, 4096), nn.ReLU(),
nn.Linear(4096, 4096), nn.ReLU())
self.cls_loc = nn.Linear(4096, n_class * 4)
self.score = nn.Linear(4096, n_class)
self.n_class = n_class
self.roi_size = roi_size
self.spatial_scale = spatial_scale
self.roi = ROIAlign((self.roi_size, self.roi_size),
self.spatial_scale,
sampling_ratio=sampling_ratio)
init.gauss_(self.cls_loc.weight, 0, 0.001)
init.constant_(self.cls_loc.bias, 0)
init.gauss_(self.score.weight, 0, 0.01)
init.constant_(self.score.bias, 0)
def initialize_weights(net):
for m in net.modules():
if isinstance(m, nn.Conv):
init.gauss_(m.weight, 0, 0.02)
init.constant_(m.bias, 0)
elif isinstance(m, nn.ConvTranspose):
init.gauss_(m.weight, 0, 0.02)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.gauss_(m.weight, 0, 0.02)
init.constant_(m.bias, 0)
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
norm_layer=nn.BatchNorm,
dilation=1,
use_dcn=False):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv(inplanes,
planes,
kernel_size=1,
bias=False,
dilation=dilation)
self.bn1 = norm_layer(planes)
if use_dcn:
self.conv2 = DCN(planes,
planes,
kernel_size=3,
stride=stride,
padding=dilation,
dilation=dilation,
deformable_groups=1)
init.constant_(self.conv2.bias, 0.0)
init.constant_(self.conv2.conv_offset_mask.weight, 0.0)
init.constant_(self.conv2.conv_offset_mask.bias, 0.0)
else:
self.conv2 = nn.Conv(planes,
planes,
kernel_size=3,
stride=stride,
padding=dilation,
bias=False,
dilation=dilation)
self.bn2 = norm_layer(planes)
self.conv3 = nn.Conv(planes,
planes * 4,
kernel_size=1,
bias=False,
dilation=dilation)
self.bn3 = norm_layer(planes * 4)
self.relu = nn.ReLU()
self.downsample = downsample
self.stride = stride
def _normal_init(self):
for var in [self.conv1, self.score, self.loc]:
init.gauss_(var.weight, 0, 0.01)
init.constant_(var.bias, 0.0)
