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, in_channels, out_channels, stride=1, bias=False):
super().__init__()
out_ch = out_channels
groups = in_channels
kernel = 3
# print(kernel, 'x', kernel, 'x', out_channels, 'x', out_channels, 'DepthWise')
self.add_module(
'dwconv',
Conv2d(groups,
groups,
kernel_size=3,
stride=stride,
padding=1,
groups=groups,
bias=bias))
self.add_module('norm', FrozenBatchNorm2d(groups))
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 __init__(self, chi, cho, norm):
super(DeformConv, self).__init__()
self.actf = nn.Sequential(get_norm(norm, cho), nn.ReLU(inplace=True))
self.offset = Conv2d(chi,
27,
kernel_size=3,
stride=1,
padding=1,
dilation=1)
self.conv = ModulatedDeformConv(chi,
cho,
kernel_size=3,
stride=1,
padding=1,
dilation=1,
deformable_groups=1)
nn.init.constant_(self.offset.weight, 0)
nn.init.constant_(self.offset.bias, 0)
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, cfg: CfgNode, input_channels: int):
"""
Initialize DensePose fully convolutional head
Args:
cfg (CfgNode): configuration options
input_channels (int): number of input channels
"""
super(DensePoseV1ConvXHead, self).__init__()
# fmt: off
hidden_dim = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL
self.n_stacked_convs = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS
self.depthwise_on = cfg.MODEL.ROI_DENSEPOSE_HEAD.DEPTHWISE.DEPTHWISE_ON
depthwise_norms = cfg.MODEL.ROI_DENSEPOSE_HEAD.DEPTHWISE.NORMS
depthwise_activations = cfg.MODEL.ROI_DENSEPOSE_HEAD.DEPTHWISE.ACTIVATIONS
# fmt: on
pad_size = kernel_size // 2
n_channels = input_channels
for i in range(self.n_stacked_convs):
if self.depthwise_on:
layer = DepthwiseSeparableConv2d(
n_channels,
hidden_dim,
kernel_size=kernel_size,
padding=pad_size,
norm1=depthwise_norms[0],
activation1=nn.ReLU()
if depthwise_activations[0] else None,
norm2=depthwise_norms[1],
activation2=nn.ReLU()
if depthwise_activations[1] else None)
else:
layer = Conv2d(n_channels,
hidden_dim,
kernel_size,
stride=1,
padding=pad_size)
layer_name = self._get_layer_name(i)
self.add_module(layer_name, layer) # pyre-ignore[16]
n_channels = hidden_dim
self.n_out_channels = n_channels
if not self.depthwise_on:
initialize_module_params(self)
def __init__(self, input_shape, *, num_keypoints, conv_dims, **kwargs):
"""
NOTE: this interface is experimental.
Args:
input_shape (ShapeSpec): shape of the input feature
conv_dims: an iterable of output channel counts for each conv in the head
e.g. (512, 512, 512) for three convs outputting 512 channels.
"""
num_keypoints = num_keypoints * 2
super().__init__(num_keypoints=num_keypoints, **kwargs)
# default up_scale to 2 (this can be made an option)
up_scale = 2
in_channels = input_shape.channels
self.blocks = []
for idx, layer_channels in enumerate(conv_dims, 1):
module = Conv2d(in_channels,
layer_channels,
3,
stride=1,
padding=1)
self.add_module("conv_fcn{}".format(idx), module)
self.blocks.append(module)
in_channels = layer_channels
deconv_kernel = 4
self.score_lowres = ConvTranspose2d(in_channels,
num_keypoints,
deconv_kernel,
stride=2,
padding=deconv_kernel // 2 - 1)
self.up_scale = up_scale
for name, param in self.named_parameters():
if "bias" in name:
nn.init.constant_(param, 0)
elif "weight" in name:
# Caffe2 implementation uses MSRAFill, which in fact
# corresponds to kaiming_normal_ in PyTorch
nn.init.kaiming_normal_(param,
mode="fan_out",
nonlinearity="relu")
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,
*,
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, cfg, input_shape: ShapeSpec):
super(MaskIoUHead, self).__init__()
# fmt: off
num_classes = cfg.MODEL.ROI_HEADS.NUM_CLASSES
conv_dims = cfg.MODEL.ROI_MASKIOU_HEAD.CONV_DIM
num_conv = cfg.MODEL.ROI_MASKIOU_HEAD.NUM_CONV
input_channels = input_shape.channels + 1
resolution = input_shape.width // 2
# fmt: on
self.conv_relus = []
stride = 1
for k in range(num_conv):
if (k + 1) == num_conv:
stride = 2
conv = Conv2d(input_channels if k == 0 else conv_dims,
conv_dims,
kernel_size=3,
stride=stride,
padding=1,
activation=F.relu)
self.add_module("maskiou_fcn{}".format(k + 1), conv)
self.conv_relus.append(conv)
self.maskiou_fc1 = Linear(conv_dims * resolution**2, 1024)
self.maskiou_fc2 = Linear(1024, 1024)
self.maskiou = Linear(1024, num_classes)
self.pooling = MaxPool2d(kernel_size=2, stride=2)
for l in self.conv_relus:
nn.init.kaiming_normal_(l.weight,
mode="fan_out",
nonlinearity="relu")
nn.init.constant_(l.bias, 0)
for l in [self.maskiou_fc1, self.maskiou_fc2]:
nn.init.kaiming_normal_(l.weight,
mode="fan_out",
nonlinearity="relu")
nn.init.constant_(l.bias, 0)
nn.init.normal_(self.maskiou.weight, mean=0, std=0.01)
nn.init.constant_(self.maskiou.bias, 0)
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 __init__(self, cfg, input_channels):
super(DensePoseV1ConvXHead, self).__init__()
# fmt: off
hidden_dim = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL
self.n_stacked_convs = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS
# fmt: on
pad_size = kernel_size // 2
n_channels = input_channels
for i in range(self.n_stacked_convs):
layer = Conv2d(n_channels,
hidden_dim,
kernel_size,
stride=1,
padding=pad_size)
layer_name = self._get_layer_name(i)
self.add_module(layer_name, layer)
n_channels = hidden_dim
self.n_out_channels = n_channels
initialize_module_params(self)
def __init__(self, cfg):
super().__init__()
in_channel = cfg.MODEL.FPN.OUT_CHANNELS
conv_dims = cfg.MODEL.KERNEL_HEAD.CONVS_DIM
num_convs = cfg.MODEL.KERNEL_HEAD.NUM_CONVS
deform = cfg.MODEL.KERNEL_HEAD.DEFORM
coord = cfg.MODEL.KERNEL_HEAD.COORD
norm = cfg.MODEL.KERNEL_HEAD.NORM
self.kernel_head = SingleHead(cfg, in_channel+2 if coord else in_channel,
conv_dims,
num_convs,
deform=deform,
coord=coord,
norm=norm,
name='kernel_head')
self.out_conv = Conv2d(conv_dims, conv_dims, kernel_size=3, padding=1)
nn.init.normal_(self.out_conv.weight, mean=0, std=0.01)
if self.out_conv.bias is not None:
nn.init.constant_(self.out_conv.bias, 0)
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, cfg, input_shape: ShapeSpec):
super(DefGridHead, self).__init__()
self.device = cfg.MODEL.DEVICE
self.grid_size = cfg.MODEL.DEFGRID_MASK_HEAD.GRID_SIZE # [20,20]
self.grid_type = cfg.MODEL.DEFGRID_MASK_HEAD.GRID_TYPE # dense_quad
self.state_dim = cfg.MODEL.DEFGRID_MASK_HEAD.STATE_DIM # 128
self.out_dim = cfg.MODEL.DEFGRID_MASK_HEAD.OUT_DIM
self.num_classes = cfg.MODEL.ROI_HEADS.NUM_CLASSES
self.sigma = cfg.MODEL.DEFGRID_MASK_HEAD.SIGMA
self.mask_coef = cfg.MODEL.DEFGRID_MASK_HEAD.MASK_COEF
self.w_variance = cfg.MODEL.DEFGRID_MASK_HEAD.W_VARIANCE
self.w_area = cfg.MODEL.DEFGRID_MASK_HEAD.W_AREA
self.w_laplacian = cfg.MODEL.DEFGRID_MASK_HEAD.W_LAPLACIAN
self.w_reconstruct_loss = cfg.MODEL.DEFGRID_MASK_HEAD.W_RECONSTRUCT_LOSS
self.matrix = MatrixUtils(1, self.grid_size, self.grid_type, self.device)
self.model = DeformableGrid(cfg, self.device)
self.to_three_channel = ConvTranspose2d(
cfg.MODEL.ROI_MASK_HEAD.CONV_DIM, 3, kernel_size=2, stride=2, padding=0
)
# self.to_three_channel = Conv2d(cfg.MODEL.ROI_MASK_HEAD.CONV_DIM, 3, kernel_size=1, stride=1, padding=0)
self.superpixel = LatticeVariance(
28,
28,
sigma=self.sigma,
device=self.device,
add_seg=True,
mask_coef=self.mask_coef,
)
self.mask_deconv = ConvTranspose2d(
self.out_dim, self.out_dim, kernel_size=2, stride=2, padding=0
)
self.mask_predictor = Conv2d(
self.out_dim, self.num_classes, kernel_size=1, stride=1, padding=0
)
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
dilation=1,
padding='',
bias=False,
norm='',
act_layer=Swish):
super(ConvBnAct2d, self).__init__()
# self.conv = create_conv2d(
# in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias)
self.conv = Conv2d(in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
padding=kernel_size // 2,
bias=(norm == ''))
self.bn = get_norm(norm, out_channels)
self.act = None if act_layer is None else act_layer(inplace=True)
def __init__(self, cfg, input_shape: ShapeSpec, **model_kwargs):
"""
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().__init__()
# fmt: off
num_classes = cfg.MODEL.ROI_HEADS.NUM_CLASSES
norm = cfg.MODEL.ROI_MASK_HEAD.NORM
input_channels = input_shape.channels
cls_agnostic_mask = cfg.MODEL.ROI_MASK_HEAD.CLS_AGNOSTIC_MASK
# fmt: on
block_args = model_kwargs.pop('block_args')
assert norm is None or norm in ["BN", "SyncBN"]
if norm == "BN":
model_kwargs['norm_layer'] = BatchNorm2d
elif norm == "SyncBN":
model_kwargs['norm_layer'] = NaiveSyncBatchNorm
builder = EfficientNetBuilder(**model_kwargs)
self.conv = nn.Sequential(*builder(input_channels, block_args))
output_channels = block_args[-1][-1]['out_chs']
num_mask_classes = 1 if cls_agnostic_mask else num_classes
self.predictor = Conv2d(output_channels,
num_mask_classes,
kernel_size=1,
stride=1,
padding=0)
initialize_weight_default(self.conv)
# 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: 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,
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,
in_channels,
out_channels,
kernel=3,
stride=1,
dropout=0.1,
bias=False):
super().__init__()
out_ch = out_channels
groups = 1
# print(kernel, 'x', kernel, 'x', in_channels, 'x', out_channels)
self.add_module(
'conv',
Conv2d(in_channels,
out_ch,
kernel_size=kernel,
stride=stride,
padding=kernel // 2,
groups=groups,
bias=bias))
self.add_module('norm', FrozenBatchNorm2d(out_ch))
self.add_module('relu', nn.ReLU6(True))
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, cfg, input_shape: ShapeSpec):
"""
The following attributes are parsed from config:
conv_dims: an iterable of output channel counts for each conv in the head
e.g. (512, 512, 512) for three convs outputting 512 channels.
num_keypoints: number of keypoint heatmaps to predicts, determines the number of
channels in the final output.
"""
super(KRCNNConvDeconvUpsampleHead, self).__init__()
# fmt: off
# default up_scale to 2 (this can eventually be moved to config)
up_scale = 2
conv_dims = cfg.MODEL.ROI_KEYPOINT_HEAD.CONV_DIMS
num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.NUM_KEYPOINTS
in_channels = input_shape.channels
# fmt: on
self.blocks = []
for idx, layer_channels in enumerate(conv_dims, 1):
module = Conv2d(in_channels, layer_channels, 3, stride=1, padding=1)
self.add_module("conv_fcn{}".format(idx), module)
self.blocks.append(module)
in_channels = layer_channels
deconv_kernel = 4
self.score_lowres = ConvTranspose2d(
in_channels, num_keypoints, deconv_kernel, stride=2, padding=deconv_kernel // 2 - 1
)
self.up_scale = up_scale
for name, param in self.named_parameters():
if "bias" in name:
nn.init.constant_(param, 0)
elif "weight" in name:
# Caffe2 implementation uses MSRAFill, which in fact
# corresponds to kaiming_normal_ in PyTorch
nn.init.kaiming_normal_(param, mode="fan_out", nonlinearity="relu")
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])
def __init__(self,
input_shape,
*,
num_classes,
conv_dims,
conv_norm="",
**kwargs):
freeze_layers = kwargs['freeze_layers']
del kwargs['freeze_layers']
super().__init__(input_shape,
num_classes=num_classes,
conv_dims=conv_dims,
conv_norm=conv_norm,
**kwargs)
self.predictor_delta = Conv2d(self.predictor.in_channels,
num_classes,
kernel_size=1,
stride=1,
padding=0)
nn.init.constant_(self.predictor_delta.weight, 0.)
if self.predictor_delta.bias is not None:
nn.init.constant_(self.predictor_delta.bias, 0.)
self._freeze_layers(freeze_layers)
def __init__(self, cfg: CfgNode, input_channels: int):
"""
Initialize DensePose fully convolutional head
Args:
cfg (CfgNode): configuration options
input_channels (int): number of input channels
"""
super(DensePoseV1ConvXHead, self).__init__()
# fmt: off
hidden_dim = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL
self.n_stacked_convs = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS
# fmt: on
pad_size = kernel_size // 2
n_channels = input_channels
for i in range(self.n_stacked_convs):
layer = Conv2d(n_channels, hidden_dim, kernel_size, stride=1, padding=pad_size)
layer_name = self._get_layer_name(i)
self.add_module(layer_name, layer)
n_channels = hidden_dim
self.n_out_channels = n_channels
initialize_module_params(self)
def __init__(self,
in_channels,
num_keypoints,
conv_dims,
loss_weight=1.0,
loss_normalizer=1.0):
super().__init__()
# default up_scale to 2.0 (this can be made an option)
up_scale = 2.0
for idx, layer_channels in enumerate(conv_dims, 1):
module = Conv2d(in_channels,
layer_channels,
3,
stride=1,
padding=1)
self.add_module("conv_fcn{}".format(idx), module)
self.add_module("conv_fcn_relu{}".format(idx), nn.ReLU())
in_channels = layer_channels
deconv_kernel = 4
self.score_lowres = ConvTranspose2d(in_channels,
num_keypoints,
deconv_kernel,
stride=2,
padding=deconv_kernel // 2 - 1)
self.up_scale = up_scale
for name, param in self.named_parameters():
if "bias" in name:
nn.init.constant_(param, 0)
elif "weight" in name:
# Caffe2 implementation uses MSRAFill, which in fact
# corresponds to kaiming_normal_ in PyTorch
nn.init.kaiming_normal_(param,
mode="fan_out",
nonlinearity="relu")
def __init__(self, input_shape, output_dim=32, output_stride=4, norm='BN'):
super().__init__()
# fmt: off
self.in_features = ["p2", "p3", "p4", "p5"]
feature_strides = {k: v.stride for k, v in input_shape.items()}
feature_channels = {k: v.channels for k, v in input_shape.items()}
# 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(output_stride)))
for k in range(head_length):
conv = Conv2d(
feature_channels[in_feature] if k == 0 else output_dim,
output_dim,
kernel_size=3,
stride=1,
padding=1,
bias=not norm,
norm=get_norm(norm, output_dim),
activation=F.relu,
)
weight_init.c2_msra_fill(conv)
head_ops.append(conv)
if feature_strides[in_feature] != output_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])
def replace_resnet(model, arch):
model = copy.deepcopy(model)
model.backbone.bottom_up = arch
in_channels = [arch.V[i].out_shape[1] for i in arch.out_layers]
in_channels = in_channels[::-1]
for i in range(len(in_channels)):
b = model.backbone.lateral_convs[i]
Fi = in_channels[i]
Fo = b.out_channels
# Copied from shrink_layer
groups = b.groups
if (groups == b.in_channels and b.in_channels == b.out_channels
and Fi == Fo):
groups = Fi
norm = None
if b.norm is not None:
norm = BatchNorm2d(Fo)
norm.weight = nn.Parameter(b.norm.weight[:Fo].clone().detach())
norm.bias = nn.Parameter(b.norm.bias[:Fo].clone().detach())
norm.running_mean = b.norm.running_mean[:Fo].clone().detach()
norm.running_var = b.norm.running_var[:Fo].clone().detach()
conv = Conv2d(Fi,
Fo,
b.kernel_size,
stride=b.stride,
padding=b.padding,
dilation=b.dilation,
groups=groups,
bias=(b.bias is not None),
norm=norm,
activation=b.activation)
conv.weight = nn.Parameter(b.weight[:Fo, :(Fi //
groups)].clone().detach())
if b.bias is not None:
conv.bias = nn.Parameter(b.bias[:Fo].clone().detach())
model.backbone.lateral_convs[i] = conv
return model
def __init__(self, input_shape, conv_dims, fc_dims, conv_norm=""):
super().__init__()
assert len(conv_dims) + len(fc_dims) > 0
self._output_size = (input_shape[1], input_shape[2], input_shape[3])
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=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):
if k == 0:
self.add_module("flatten", nn.Flatten())
fc = Linear(int(np.prod(self._output_size)), fc_dim)
self.add_module("fc{}".format(k + 1), fc)
self.add_module("fc_relu{}".format(k + 1), nn.ReLU())
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)
