def init_weights(self, pretrained=None):
if isinstance(pretrained, str):
logger = logging.getLogger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.modules():
if isinstance(m, nn.Conv2d):
kaiming_init(m)
elif isinstance(m, (_BatchNorm, nn.GroupNorm)):
constant_init(m, 1)
if self.dcn is not None:
for m in self.modules():
if isinstance(m, Bottleneck) and hasattr(
m, "conv2_offset"):
constant_init(m.conv2_offset, 0)
if self.zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
constant_init(m.norm3, 0)
elif isinstance(m, BasicBlock):
constant_init(m.norm2, 0)
else:
raise TypeError("pretrained must be a str or None")
def __init__(
self,
in_channels,
heads,
head_conv=64,
final_kernel=1,
bn=False,
init_bias=-2.19,
**kwargs,
):
super(SepHead, self).__init__(**kwargs)
self.heads = heads
for head in self.heads:
classes, num_conv = self.heads[head]
fc = Sequential()
for i in range(num_conv - 1):
fc.add(
nn.Conv2d(in_channels,
head_conv,
kernel_size=final_kernel,
stride=1,
padding=final_kernel // 2,
bias=True))
if bn:
fc.add(nn.BatchNorm2d(head_conv))
fc.add(nn.ReLU())
fc.add(
nn.Conv2d(head_conv,
classes,
kernel_size=final_kernel,
stride=1,
padding=final_kernel // 2,
bias=True))
if 'hm' in head:
fc[-1].bias.data.fill_(init_bias)
else:
for m in fc.modules():
if isinstance(m, nn.Conv2d):
kaiming_init(m)
self.__setattr__(head, fc)
def init_weights(self, pretrained=None):
if isinstance(pretrained, str):
logger = logging.getLogger()
load_checkpoint(self, pretrained, strict=False, logger=logger)
elif pretrained is None:
for m in self.features.modules():
if isinstance(m, nn.Conv2d):
kaiming_init(m)
elif isinstance(m, nn.BatchNorm2d):
constant_init(m, 1)
elif isinstance(m, nn.Linear):
normal_init(m, std=0.01)
else:
raise TypeError("pretrained must be a str or None")
for m in self.extra.modules():
if isinstance(m, nn.Conv2d):
xavier_init(m, distribution="uniform")
constant_init(self.l2_norm, self.l2_norm.scale)
def init_weights(self):
nonlinearity = "relu" if self.activation is None else self.activation
kaiming_init(self.conv, nonlinearity=nonlinearity)
if self.with_norm:
constant_init(self.norm, 1, bias=0)
