def __init__(self, in_channels, out_channels, *, stride=1, norm="BN"):
"""
Args:
in_channels (int): Number of input channels.
out_channels (int): Number of output channels.
stride (int): Stride for the first conv.
norm (str or callable): normalization for all conv layers.
See :func:`layers.get_norm` for supported format.
"""
super().__init__(in_channels, out_channels, stride)
if in_channels != out_channels:
self.shortcut = Conv2d(
in_channels,
out_channels,
kernel_size=1,
stride=stride,
bias=False,
norm=get_norm(norm, out_channels),
)
else:
self.shortcut = None
self.conv1 = Conv2d(
in_channels,
out_channels,
kernel_size=3,
stride=stride,
padding=1,
bias=False,
norm=get_norm(norm, out_channels),
)
self.conv2 = Conv2d(
out_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False,
norm=get_norm(norm, out_channels),
)
for layer in [self.conv1, self.conv2, self.shortcut]:
if layer is not None: # shortcut can be None
weight_init.c2_msra_fill(layer)
def __init__(self, cfg, input_shape: ShapeSpec):
"""
The following attributes are parsed from config:
num_conv, num_fc: the number of conv/fc layers
conv_dim/fc_dim: the dimension of the conv/fc layers
norm: normalization for the conv layers
"""
super().__init__()
# fmt: off
num_conv = cfg.MODEL.ROI_BOX_HEAD.NUM_CONV
conv_dim = cfg.MODEL.ROI_BOX_HEAD.CONV_DIM
num_fc = cfg.MODEL.ROI_BOX_HEAD.NUM_FC
fc_dim = cfg.MODEL.ROI_BOX_HEAD.FC_DIM
norm = cfg.MODEL.ROI_BOX_HEAD.NORM
# fmt: on
assert num_conv + num_fc > 0
self._output_size = (input_shape.channels, input_shape.height, input_shape.width)
self.conv_norm_relus = []
for k in range(num_conv):
conv = Conv2d(
self._output_size[0],
conv_dim,
kernel_size=3,
padding=1,
bias=not norm,
norm=get_norm(norm, conv_dim),
activation=F.relu,
)
self.add_module("conv{}".format(k + 1), conv)
self.conv_norm_relus.append(conv)
self._output_size = (conv_dim, self._output_size[1], self._output_size[2])
self.fcs = []
for k in range(num_fc):
fc = nn.Linear(np.prod(self._output_size), fc_dim)
self.add_module("fc{}".format(k + 1), fc)
self.fcs.append(fc)
self._output_size = fc_dim
for layer in self.conv_norm_relus:
weight_init.c2_msra_fill(layer)
for layer in self.fcs:
weight_init.c2_xavier_fill(layer)
def __init__(self, cfg, input_shape: ShapeSpec):
"""
The following attributes are parsed from config:
fc_dim: the output dimension of each FC layers
num_fc: the number of FC layers
coarse_pred_each_layer: if True, coarse prediction features are concatenated to each
layer's input
"""
super(StandardPointHead, self).__init__()
# fmt: off
num_classes = cfg.MODEL.POINT_HEAD.NUM_CLASSES
fc_dim = cfg.MODEL.POINT_HEAD.FC_DIM
num_fc = cfg.MODEL.POINT_HEAD.NUM_FC
cls_agnostic_mask = cfg.MODEL.POINT_HEAD.CLS_AGNOSTIC_MASK
self.coarse_pred_each_layer = cfg.MODEL.POINT_HEAD.COARSE_PRED_EACH_LAYER
input_channels = input_shape.channels
# fmt: on
fc_dim_in = input_channels + num_classes
self.fc_layers = []
for k in range(num_fc):
fc = nn.Conv1d(fc_dim_in,
fc_dim,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.add_module("fc{}".format(k + 1), fc)
self.fc_layers.append(fc)
fc_dim_in = fc_dim
fc_dim_in += num_classes if self.coarse_pred_each_layer else 0
num_mask_classes = 1 if cls_agnostic_mask else num_classes
self.predictor = nn.Conv1d(fc_dim_in,
num_mask_classes,
kernel_size=1,
stride=1,
padding=0)
for layer in self.fc_layers:
weight_init.c2_msra_fill(layer)
# use normal distribution initialization for mask prediction layer
nn.init.normal_(self.predictor.weight, std=0.001)
if self.predictor.bias is not None:
nn.init.constant_(self.predictor.bias, 0)
def __init__(self, input_channels, num_classes):
"""
The following attributes are parsed from config:
fc_dim: the output dimension of each FC layers
num_fc: the number of FC layers
coarse_pred_each_layer: if True, coarse prediction features are concatenated to each
layer's input
"""
super(StandardPointHead, self).__init__()
# fmt: off
num_classes = num_classes
fc_dim = 256
num_fc = 3
cls_agnostic_mask = False
self.coarse_pred_each_layer = True
input_channels = input_channels
# fmt: on
fc_dim_in = input_channels + num_classes
self.fc_layers = []
for k in range(num_fc):
fc = nn.Conv1d(fc_dim_in,
fc_dim,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.add_module("fc{}".format(k + 1), fc)
self.fc_layers.append(fc)
fc_dim_in = fc_dim
fc_dim_in += num_classes if self.coarse_pred_each_layer else 0
num_mask_classes = 1 if cls_agnostic_mask else num_classes
self.predictor = nn.Conv1d(fc_dim_in,
num_mask_classes,
kernel_size=1,
stride=1,
padding=0)
for layer in self.fc_layers:
weight_init.c2_msra_fill(layer)
# use normal distribution initialization for mask prediction layer
nn.init.normal_(self.predictor.weight, std=0.001)
if self.predictor.bias is not None:
nn.init.constant_(self.predictor.bias, 0)
def init_weights(
model, fc_init_std=0.01, zero_init_final_bn=True, zero_init_final_conv=False
):
"""
Performs ResNet style weight initialization.
Args:
fc_init_std (float): the expected standard deviation for fc layer.
zero_init_final_bn (bool): if True, zero initialize the final bn for
every bottleneck.
"""
for m in model.modules():
if isinstance(m, nn.Conv3d):
# Note that there is no bias due to BN
if hasattr(m, "final_conv") and zero_init_final_conv:
m.weight.data.zero_()
else:
"""
Follow the initialization method proposed in:
{He, Kaiming, et al.
"Delving deep into rectifiers: Surpassing human-level
performance on imagenet classification."
arXiv preprint arXiv:1502.01852 (2015)}
"""
c2_msra_fill(m)
elif isinstance(m, (nn.BatchNorm3d, nn.BatchNorm2d, nn.BatchNorm1d)):
if (
hasattr(m, "transform_final_bn")
and m.transform_final_bn
and zero_init_final_bn
):
batchnorm_weight = 0.0
else:
batchnorm_weight = 1.0
if m.weight is not None:
m.weight.data.fill_(batchnorm_weight)
if m.bias is not None:
m.bias.data.zero_()
if isinstance(m, nn.Linear):
if hasattr(m, "xavier_init") and m.xavier_init:
c2_xavier_fill(m)
else:
m.weight.data.normal_(mean=0.0, std=fc_init_std)
if m.bias is not None:
m.bias.data.zero_()
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
super().__init__()
# fmt: off
self.in_features = cfg.MODEL.SEG_DET_HEAD.IN_FEATURES
feature_strides = {k: v.stride for k, v in input_shape.items()}
feature_channels = {k: v.channels for k, v in input_shape.items()}
self.ignore_value = cfg.MODEL.SEG_DET_HEAD.IGNORE_VALUE
num_classes = cfg.MODEL.SEG_DET_HEAD.NUM_CLASSES
conv_dims = cfg.MODEL.SEG_DET_HEAD.CONVS_DIM
self.common_stride = cfg.MODEL.SEG_DET_HEAD.COMMON_STRIDE
norm = cfg.MODEL.SEG_DET_HEAD.NORM
self.loss_weight = cfg.MODEL.SEG_DET_HEAD.LOSS_WEIGHT
# fmt: on
self.scale_heads = []
for in_feature in self.in_features:
head_ops = []
head_length = max(
1, int(np.log2(feature_strides[in_feature]) - np.log2(self.common_stride))
)
for k in range(head_length):
norm_module = nn.GroupNorm(32, conv_dims) if norm == "GN" else None
conv = Conv2d(
feature_channels[in_feature] if k == 0 else conv_dims,
conv_dims,
kernel_size=3,
stride=1,
padding=1,
bias=not norm,
norm=norm_module,
activation=F.relu,
)
weight_init.c2_msra_fill(conv)
head_ops.append(conv)
if feature_strides[in_feature] != self.common_stride:
head_ops.append(
nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
)
self.scale_heads.append(nn.Sequential(*head_ops))
self.add_module(in_feature, self.scale_heads[-1])
self.predictor_segmap = Conv2d(conv_dims, num_classes, kernel_size=1, stride=1, padding=0)
self.predictor_contour = Conv2d(conv_dims, num_classes, kernel_size=1, stride=1, padding=0)
# Embedding to separate different objects of the same category
self.predictor_emb = Conv2d(conv_dims, num_classes, kernel_size=1, stride=1, padding=0)
weight_init.c2_msra_fill(self.predictor_segmap)
weight_init.c2_msra_fill(self.predictor_contour)
weight_init.c2_msra_fill(self.predictor_emb)
self.emb_loss = Embedding_loss(cfg)
self.device = torch.device(cfg.MODEL.DEVICE)
self.num_classes = num_classes
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
"""
Args:
norm (str or callable): a callable that takes the number of
channels and return a `nn.Module`, or a pre-defined string
(one of {"FrozenBN", "BN", "GN"}).
"""
super().__init__()
self.conv1 = Conv2d(
in_channels,
out_channels,
kernel_size=7,
stride=2,
padding=3,
bias=False,
norm=get_norm(norm, out_channels),
)
weight_init.c2_msra_fill(self.conv1)
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
"""
Args:
norm (str or callable): norm after the first conv layer.
See :func:`layers.get_norm` for supported format.
"""
super().__init__(in_channels, out_channels, 4)
self.in_channels = in_channels
self.conv1 = Conv2d(
in_channels,
out_channels,
kernel_size=7,
stride=2,
padding=3,
bias=False,
norm=get_norm(norm, out_channels),
)
weight_init.c2_msra_fill(self.conv1)
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
super().__init__()
num_classes = cfg.MODEL.SEMANTIC_INSTANCE_HEAD.NUM_CLASSES
feature_channels = input_shape.channels
self.huber_active = cfg.MODEL.SEMANTIC_INSTANCE_HEAD.HUBER_ACTIVE
self.dice_active = cfg.MODEL.SEMANTIC_INSTANCE_HEAD.DICE_ACTIVE
self.loss_weight = cfg.MODEL.SEMANTIC_INSTANCE_HEAD.LOSS_WEIGHT
self.ignore_value = cfg.MODEL.SEMANTIC_INSTANCE_HEAD.IGNORE_VALUE
self.dual_loss = cfg.MODEL.SEMANTIC_INSTANCE_HEAD.DUAL_LOSS
if self.dual_loss:
self.dual_loss_weight = \
cfg.MODEL.SEMANTIC_INSTANCE_HEAD.DUAL_LOSS_WEIGHT
self.predictor = Conv2d(feature_channels,
num_classes,
kernel_size=1,
stride=1,
padding=0)
weight_init.c2_msra_fill(self.predictor)
def test_conv_weight_init(self):
# Test weight initialization for convolutional layers.
kernel_sizes = [1, 3]
channel_in_dims = [128, 256, 512]
channel_out_dims = [256, 512, 1024]
for layer in [nn.Conv1d, nn.Conv2d, nn.Conv3d]:
for k_size, c_in_dim, c_out_dim in itertools.product(
kernel_sizes, channel_in_dims, channel_out_dims):
p = {
"kernel_size": k_size,
"in_channels": c_in_dim,
"out_channels": c_out_dim,
}
model = layer(**p)
if layer is nn.Conv1d:
spatial_dim = k_size
elif layer is nn.Conv2d:
spatial_dim = k_size**2
elif layer is nn.Conv3d:
spatial_dim = k_size**3
# Calculate fan_in and fan_out.
fan_in = c_in_dim * spatial_dim
fan_out = c_out_dim * spatial_dim
# Msra weight init check.
c2_msra_fill(model)
self.assertTrue(
TestWeightInit.weight_and_bias_dist_match(
model.weight,
model.bias,
TestWeightInit.msra_fill_std(fan_out),
))
# Xavier weight init check.
c2_xavier_fill(model)
self.assertTrue(
TestWeightInit.weight_and_bias_dist_match(
model.weight,
model.bias,
TestWeightInit.xavier_fill_std(fan_in),
))
def __init__(
self,
gate_channel: int,
reduction_ratio: int = 16,
dilation_conv_num: int = 2,
dilation_val: int = 4,
) -> None:
super().__init__()
self.gate_s = nn.Sequential()
self.gate_s.add_module(
"gate_s_conv_reduce0",
nn.Conv2d(gate_channel,
gate_channel // reduction_ratio,
kernel_size=1),
)
self.gate_s.add_module(
"gate_s_bn_reduce0",
nn.BatchNorm2d(gate_channel // reduction_ratio),
)
self.gate_s.add_module("gate_s_relu_reduce0", nn.ReLU())
for i in range(dilation_conv_num):
self.gate_s.add_module(
"gate_s_conv_di_%d" % i,
nn.Conv2d(
gate_channel // reduction_ratio,
gate_channel // reduction_ratio,
kernel_size=3,
padding=dilation_val,
dilation=dilation_val,
),
)
self.gate_s.add_module(
"gate_s_bn_di_%d" % i,
nn.BatchNorm2d(gate_channel // reduction_ratio),
)
self.gate_s.add_module("gate_s_relu_di_%d" % i, nn.ReLU())
self.gate_s.add_module(
"gate_s_conv_final",
nn.Conv2d(gate_channel // reduction_ratio, 1, kernel_size=1),
)
for m in self.modules():
if type(m) == nn.Conv2d:
weight_init.c2_msra_fill(m)
def __init__(
self, input_shape: ShapeSpec, *, conv_dims: List[int], fc_dims: List[int], conv_norm=""
):
"""
NOTE: this interface is experimental.
Args:
input_shape (ShapeSpec): shape of the input feature.
conv_dims (list[int]): the output dimensions of the conv layers
fc_dims (list[int]): the output dimensions of the fc layers
conv_norm (str or callable): normalization for the conv layers.
See :func:`detectron2.layers.get_norm` for supported types.
"""
super().__init__()
assert len(conv_dims) + len(fc_dims) > 0
self._output_size = (input_shape.channels, input_shape.height, input_shape.width)
self.conv_norm_relus = []
for k, conv_dim in enumerate(conv_dims):
conv = Conv2d(
self._output_size[0],
conv_dim,
kernel_size=3,
padding=1,
bias=not conv_norm,
norm=get_norm(conv_norm, conv_dim),
activation=F.relu,
)
self.add_module("conv{}".format(k + 1), conv)
self.conv_norm_relus.append(conv)
self._output_size = (conv_dim, self._output_size[1], self._output_size[2])
self.fcs = []
for k, fc_dim in enumerate(fc_dims):
fc = Linear(np.prod(self._output_size), fc_dim)
self.add_module("fc{}".format(k + 1), fc)
self.fcs.append(fc)
self._output_size = fc_dim
for layer in self.conv_norm_relus:
weight_init.c2_msra_fill(layer)
for layer in self.fcs:
weight_init.c2_xavier_fill(layer)
def __init__(
self,
in_channels,
out_channels,
kernel_size=3,
padding=1,
dilation=1,
*,
norm1=None,
activation1=None,
norm2=None,
activation2=None,
):
"""
Args:
norm1, norm2 (str or callable): normalization for the two conv layers.
activation1, activation2 (callable(Tensor) -> Tensor): activation
function for the two conv layers.
"""
super().__init__()
self.depthwise = Conv2d(
in_channels,
in_channels,
kernel_size=kernel_size,
padding=padding,
dilation=dilation,
groups=in_channels,
bias=not norm1,
norm=get_norm(norm1, in_channels),
activation=activation1,
)
self.pointwise = Conv2d(
in_channels,
out_channels,
kernel_size=1,
bias=not norm2,
norm=get_norm(norm2, out_channels),
activation=activation2,
)
# default initialization
weight_init.c2_msra_fill(self.depthwise)
weight_init.c2_msra_fill(self.pointwise)
def __init__(
self,
in_channels,
out_channels,
*,
bottleneck_channels,
stride=1,
num_groups=1,
norm="BN",
stride_in_1x1=False,
dilation=1,
):
"""
Args:
bottleneck_channels (int): number of output channels for the 3x3
"bottleneck" conv layers.
num_groups (int): number of groups for the 3x3 conv layer.
norm (str or callable): normalization for all conv layers.
See :func:`layers.get_norm` for supported format.
stride_in_1x1 (bool): when stride>1, whether to put stride in the
first 1x1 convolution or the bottleneck 3x3 convolution.
dilation (int): the dilation rate of the 3x3 conv layer.
"""
super().__init__(in_channels, out_channels, stride)
self.shortcut = None
# 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=(2,2),
padding=1,
bias=False,
norm=get_norm(norm, bottleneck_channels),
)
weight_init.c2_msra_fill(self.conv1)
def __init__(self, in_channels=3, out_channels=64, norm="BN"):
"""
Args:
norm (str or callable): norm after the first conv layer.
See :func:`layers.get_norm` for supported format.
"""
super().__init__(in_channels, out_channels, 4)
self.in_channels = in_channels
self.conv1 = nn.Sequential(
Conv2d(
in_channels,
32,
kernel_size=3,
stride=2,
padding=1,
bias=False,
),
get_norm(norm, 32),
nn.ReLU(inplace=True),
Conv2d(
32,
32,
kernel_size=3,
stride=1,
padding=1,
bias=False,
),
get_norm(norm, 32),
nn.ReLU(inplace=True),
Conv2d(
32,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False,
),
)
self.bn1 = get_norm(norm, out_channels)
for layer in self.conv1:
if isinstance(layer, Conv2d):
weight_init.c2_msra_fill(layer)
def __init__(self, input_shape: ShapeSpec, *, num_classes, conv_dims, conv_norm="", **kwargs):
"""
NOTE: this interface is experimental.
Args:
input_shape (ShapeSpec): shape of the input feature
num_classes (int): the number of classes. 1 if using class agnostic prediction.
conv_dims (list[int]): a list of N>0 integers representing the output dimensions
of N-1 conv layers and the last upsample layer.
conv_norm (str or callable): normalization for the conv layers.
See :func:`detectron2.layers.get_norm` for supported types.
"""
super().__init__(**kwargs)
assert len(conv_dims) >= 1, "conv_dims have to be non-empty!"
self.conv_norm_relus = []
cur_channels = input_shape.channels
for k, conv_dim in enumerate(conv_dims[:-1]):
conv = Conv2d(
cur_channels,
conv_dim,
kernel_size=3,
stride=1,
padding=1,
bias=not conv_norm,
norm=get_norm(conv_norm, conv_dim),
activation=nn.ReLU(),
)
self.add_module("mask_fcn{}".format(k + 1), conv)
self.conv_norm_relus.append(conv)
cur_channels = conv_dim
self.deconv = ConvTranspose2d(
cur_channels, conv_dims[-1], kernel_size=2, stride=2, padding=0
)
self.add_module("deconv_relu", nn.ReLU())
cur_channels = conv_dims[-1]
for layer in self.conv_norm_relus + [self.deconv]:
weight_init.c2_msra_fill(layer)
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
super().__init__()
# fmt: off
self.in_features = cfg.MODEL.FPN_BLOCK.IN_FEATURES
feature_strides = {k: v.stride for k, v in input_shape.items()}
feature_channels = {k: v.channels for k, v in input_shape.items()}
self.conv_dims = cfg.MODEL.FPN_BLOCK.CONVS_DIM
self.common_stride = cfg.MODEL.FPN_BLOCK.COMMON_STRIDE
norm = cfg.MODEL.FPN_BLOCK.NORM
# fmt: on
self.scale_blocks = []
for in_feature in self.in_features:
block_ops = []
block_length = max(
1,
int(
np.log2(feature_strides[in_feature]) -
np.log2(self.common_stride)))
for k in range(block_length):
norm_module = nn.GroupNorm(
32, self.conv_dims) if norm == "GN" else None
conv = Conv2d(
feature_channels[in_feature] if k == 0 else self.conv_dims,
self.conv_dims,
kernel_size=3,
stride=1,
padding=1,
bias=not norm_module,
norm=norm_module,
activation=F.relu,
)
weight_init.c2_msra_fill(conv)
block_ops.append(conv)
if feature_strides[in_feature] != self.common_stride:
block_ops.append(
nn.Upsample(scale_factor=2,
mode="bilinear",
align_corners=True))
self.scale_blocks.append(nn.Sequential(*block_ops))
self.add_module(in_feature, self.scale_blocks[-1])
def conv_dw(inp, oup, stride, width_mult=1.0):
inp = 3 if inp == 3 else scale_chan(inp, width_mult)
oup = scale_chan(oup, width_mult)
conv1 = Conv2d(inp,
inp,
3,
stride,
1,
groups=inp,
bias=False,
norm=get_norm("BN", inp))
conv2 = Conv2d(inp, oup, 1, 1, 0, bias=False, norm=get_norm("BN", oup))
weight_init.c2_msra_fill(conv1)
weight_init.c2_msra_fill(conv2)
return nn.Sequential(
conv1,
nn.ReLU(inplace=True),
conv2,
nn.ReLU(inplace=True),
)
def __init__(self, cfg, input_shape: ShapeSpec):
"""
The following attributes are parsed from config:
num_conv: the number of conv layers
conv_dim: the dimension of the conv layers
norm: normalization for the conv layers
"""
super(GeneralDis, self).__init__()
# fmt: off
assert input_shape.height == input_shape.width
num_resblock = cfg.MODEL.ROI_DIS_HEAD.NUM_RESBLOCK
conv_dims = cfg.MODEL.ROI_DIS_HEAD.CONV_DIM
self.norm = cfg.MODEL.ROI_DIS_HEAD.NORM
# fmt: on
self.net = []
for k in range(num_resblock):
conv = Conv2d(
1 if k == 0 else conv_dims,
conv_dims,
kernel_size=3,
stride=2,
padding=1,
bias=not self.norm,
norm=get_norm(self.norm, conv_dims, num_groups=2),
activation=F.relu,
)
self.add_module("dis_conv{}".format(k + 1), conv)
self.net.append(conv)
res_block = ResBlock(conv_dims, norm=self.norm, num_group=2)
self.add_module("dis_res_block{}".format(k + 1), res_block)
self.net.append(res_block)
for layer in self.net:
if isinstance(layer, Conv2d):
weight_init.c2_msra_fill(layer)
self.aver_pooling = nn.AvgPool2d(kernel_size=4)
self.cls = nn.Linear(conv_dims, 1)
def __init__(self, in_channels, out_channels):
super().__init__()
self.num_levels = 2
self.in_feature = "p5"
norm = "GN"
conv_fcn = []
conv_fcn.append(Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=1, bias=not norm,\
groups=in_channels, norm=get_norm(norm, in_channels), activation=F.relu))
conv_fcn.append(Conv2d(in_channels, out_channels, kernel_size=1, bias=not norm,\
norm=get_norm(norm, out_channels), activation=F.relu))
self.add_module('p6', nn.Sequential(*conv_fcn))
conv_fcn = []
conv_fcn.append(Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1, bias=not norm,\
groups=out_channels, norm=get_norm(norm, out_channels), activation=F.relu))
conv_fcn.append(Conv2d(out_channels, out_channels, kernel_size=1, bias=not norm,\
norm=get_norm(norm, out_channels), activation=F.relu))
self.add_module('p7', nn.Sequential(*conv_fcn))
for layer in [*self.p6, *self.p7]:
weight_init.c2_msra_fill(layer)
def __init__(self, cfg):
super().__init__()
# fmt: off
in_features = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECODER_OUT_DIMS
num_classes = cfg.MODEL.ROI_HEADS.NUM_CLASSES
self.loss_weight = cfg.MODEL.ROI_DENSEPOSE_HEAD.SEMSEG_WEIGHTS
# fmt: on
self.sem_projector = Conv2d(in_features,
in_features,
kernel_size=3,
stride=1,
padding=1)
self.predictor = Conv2d(in_features,
num_classes + 1,
kernel_size=1,
stride=1,
padding=0)
weight_init.c2_msra_fill(self.sem_projector)
weight_init.c2_msra_fill(self.predictor)
def __init__(self, cfg: CfgNode, input_channels: int):
super(DensePoseDeepLabHead, self).__init__()
# fmt: off
hidden_dim = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL
norm = cfg.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NORM
self.n_stacked_convs = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS
self.use_nonlocal = cfg.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NONLOCAL_ON
# fmt: on
pad_size = kernel_size // 2
n_channels = input_channels
self.ASPP = ASPP(input_channels, [6, 12, 56], n_channels) # 6, 12, 56
# pyre-fixme[29]: `Union[nn.Module, torch.Tensor]` is not a function.
self.add_module("ASPP", self.ASPP)
if self.use_nonlocal:
self.NLBlock = NONLocalBlock2D(input_channels, bn_layer=True)
# pyre-fixme[29]: `Union[nn.Module, torch.Tensor]` is not a function.
self.add_module("NLBlock", self.NLBlock)
# weight_init.c2_msra_fill(self.ASPP)
for i in range(self.n_stacked_convs):
norm_module = nn.GroupNorm(32,
hidden_dim) if norm == "GN" else None
layer = Conv2d(
n_channels,
hidden_dim,
kernel_size,
stride=1,
padding=pad_size,
bias=not norm,
norm=norm_module,
)
weight_init.c2_msra_fill(layer)
n_channels = hidden_dim
layer_name = self._get_layer_name(i)
# pyre-fixme[29]: `Union[nn.Module, torch.Tensor]` is not a function.
self.add_module(layer_name, layer)
self.n_out_channels = hidden_dim
def __init__(self, cur_channels, num_classes, conv_dims, conv_norm=""):
super().__init__()
assert len(conv_dims) >= 1, "conv_dims have to be non-empty!"
self.conv_norm_relus = []
for k, conv_dim in enumerate(conv_dims[:-1]):
conv = Conv2d(
cur_channels,
conv_dim,
kernel_size=3,
stride=1,
padding=1,
bias=not conv_norm,
norm=get_norm(conv_norm, conv_dim),
activation=nn.ReLU(),
)
self.add_module("mask_fcn{}".format(k + 1), conv)
self.conv_norm_relus.append(conv)
cur_channels = conv_dim
self.deconv = ConvTranspose2d(cur_channels,
conv_dims[-1],
kernel_size=2,
stride=2,
padding=0)
self.add_module("deconv_relu", nn.ReLU())
cur_channels = conv_dims[-1]
self.predictor = Conv2d(cur_channels,
num_classes,
kernel_size=1,
stride=1,
padding=0)
for layer in self.conv_norm_relus + [self.deconv]:
weight_init.c2_msra_fill(layer)
# use normal distribution initialization for mask prediction layer
nn.init.normal_(self.predictor.weight, std=0.001)
if self.predictor.bias is not None:
nn.init.constant_(self.predictor.bias, 0)
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec], in_features):
super(Decoder, self).__init__()
# fmt: off
self.in_features = in_features
feature_strides = {k: v.stride for k, v in input_shape.items()}
feature_channels = {k: v.channels for k, v in input_shape.items()}
num_classes = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NUM_CLASSES
conv_dims = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECODER_CONV_DIMS
self.common_stride = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECODER_COMMON_STRIDE
norm = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NORM
# fmt: on
self.scale_heads = []
for in_feature in self.in_features:
head_ops = []
head_length = max(
1, int(np.log2(feature_strides[in_feature]) - np.log2(self.common_stride))
)
for k in range(head_length):
conv = Conv2d(
feature_channels[in_feature] if k == 0 else conv_dims,
conv_dims,
kernel_size=3,
stride=1,
padding=1,
bias=not norm,
norm=get_norm(norm, conv_dims),
activation=F.relu,
)
weight_init.c2_msra_fill(conv)
head_ops.append(conv)
if feature_strides[in_feature] != self.common_stride:
head_ops.append(
nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
)
self.scale_heads.append(nn.Sequential(*head_ops))
self.add_module(in_feature, self.scale_heads[-1])
self.predictor = Conv2d(conv_dims, num_classes, kernel_size=1, stride=1, padding=0)
weight_init.c2_msra_fill(self.predictor)
def __init__(
self,
in_planes,
out_planes,
deconv_kernel,
deconv_stride=2,
deconv_pad=1,
deconv_out_pad=0,
num_groups=1,
dilation=1,
modulate_deform=True # not used
):
super(CNDeconvLayer, self).__init__()
self.conv = Conv2d(
in_planes,
out_planes,
kernel_size=3,
stride=1,
padding=1,
dilation=dilation,
)
for layer in [self.conv]:
weight_init.c2_msra_fill(layer)
self.bn = nn.BatchNorm2d(out_planes)
self.up_sample = nn.ConvTranspose2d(
in_channels=out_planes,
out_channels=out_planes,
kernel_size=deconv_kernel,
stride=deconv_stride,
padding=deconv_pad,
output_padding=deconv_out_pad,
bias=False,
)
self._deconv_init()
self.up_bn = nn.BatchNorm2d(out_planes)
self.relu = nn.ReLU()
def __init__(self,
input_shape: ShapeSpec,
*,
conv_dims: List[int],
fc_dims: List[int],
conv_norm=""):
super().__init__()
assert len(conv_dims) + len(fc_dims) > 0
self._output_size = (input_shape.channels, input_shape.height,
input_shape.width)
self.conv_norm_relus = []
for k, conv_dim in enumerate(conv_dims):
conv = nn.Sequential(
nn.Conv2d(self._output_size[0],
conv_dim,
kernel_size=3,
padding=1,
bias=False), nn.BatchNorm2d(conv_dim), nn.ReLU())
self.add_module("conv{}".format(k + 1), conv)
self.conv_norm_relus.append(conv)
self._output_size = (conv_dim, self._output_size[1],
self._output_size[2])
self.fcs = []
for k, fc_dim in enumerate(fc_dims):
fc = nn.Linear(np.prod(self._output_size), fc_dim)
self.add_module("fc{}".format(k + 1), fc)
self.fcs.append(fc)
self._output_size = fc_dim
for layer in self.conv_norm_relus:
for l in layer:
if isinstance(l, nn.Conv2d):
weight_init.c2_msra_fill(l)
for layer in self.fcs:
weight_init.c2_xavier_fill(layer)
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
super().__init__()
# fmt: off
self.ignore_value = cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE
num_classes = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES
norm = cfg.MODEL.SEM_SEG_HEAD.NORM
iteration_num = cfg.MODEL.SEM_SEG_HEAD.ITERATION_NUM
em_mom = cfg.MODEL.EMA.EM_MOM
bn_mom = cfg.MODEL.EMA.BN_MOM
# fmt: on
self.reduced_conv = Conv2d(2048, 512, kernel_size=3, stride=1, padding=1, bias=False,
norm=get_norm(norm, 512, bn_mom), activation=F.relu)
self.emau = EMAUnit(512, 64, iteration_num, em_mom, norm, bn_mom)
self.predictor = nn.Sequential(
Conv2d(512, 256, kernel_size=3, stride=1, padding=1, bias=False, norm=get_norm(norm, 256, bn_mom),
activation=F.relu), nn.Dropout2d(p=0.1), Conv2d(256, num_classes, kernel_size=1))
weight_init.c2_msra_fill(self.reduced_conv)
for module in self.predictor:
if isinstance(module, Conv2d):
weight_init.c2_msra_fill(module)
def __init__(self):
super(HNet, self).__init__()
self._out_feature_strides = {"stem": 1, "s1": 1}
self._out_feature_channels = {"stem": 64, "s1": 64}
self._out_features = ["s1"]
blocks = []
self.conv1 = Conv2d(
3,
64,
kernel_size=3,
stride=1,
padding=1,
bias=False,
norm=get_norm("BN", 64),
)
weight_init.c2_msra_fill(self.conv1)
for i in range(16):
blocks.append(BottleneckBlock(64, 64, bottleneck_channels=32))
self.s1 = nn.Sequential(*blocks)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.linear = nn.Linear(64, 1000)
nn.init.normal_(self.linear.weight, std=0.01)
def __init__(self,
in_channels,
out_channels,
*,
kernel_size=3,
stride=1,
padding=0,
norm="BN"):
super(Conv2dBNLeakyReLU, self).__init__(
in_channels, out_channels, stride)
self.conv1 = nn.Conv2d(
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=False,
)
self.bn1 = get_norm(norm, out_channels)
# parameter init
weight_init.c2_msra_fill(self.conv1)
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
super().__init__()
# fmt: off
self.in_features = cfg.MODEL.SEMANTIC_FPN.IN_FEATURES
feature_strides = {k: v.stride for k, v in input_shape.items()}
feature_channels = {k: v.channels for k, v in input_shape.items()}
self.common_stride = cfg.MODEL.SEMANTIC_FPN.COMMON_STRIDE
conv_dims = cfg.MODEL.SEMANTIC_FPN.CONVS_DIM
norm = cfg.MODEL.SEMANTIC_FPN.NORM
# fmt: on
self.scale_heads = []
for in_feature in self.in_features:
head_ops = []
head_length = max(
1, int(np.log2(feature_strides[in_feature]) - np.log2(self.common_stride))
)
for k in range(head_length):
norm_module = get_norm(norm, conv_dims)
conv = Conv2d(
feature_channels[in_feature] if k == 0 else conv_dims,
conv_dims,
kernel_size=3,
stride=1,
padding=1,
bias=not norm,
norm=norm_module,
activation=F.relu,
)
weight_init.c2_msra_fill(conv)
head_ops.append(conv)
if feature_strides[in_feature] != self.common_stride:
head_ops.append(
nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
)
self.scale_heads.append(nn.Sequential(*head_ops))
self.add_module(in_feature, self.scale_heads[-1])
