def __init__(self, n_joints, shrink_width=True):
super().__init__()
def f(a, b):
if shrink_width:
return (a, b)
else:
return (b, a)
self.down_convs = nn.Sequential(
# 128 x 32 x 32
self._DownBlock(128,
256,
stride=f(1, 2),
dilation=f(2, 1),
dilation_in=f(1, 1)),
self._DownBlock(256, 256, dilation=f(2, 1)),
# 256 x 32 x 16
self._DownBlock(256,
512,
stride=f(1, 2),
dilation=f(4, 1),
dilation_in=f(2, 1)),
self._DownBlock(512, 512, dilation=f(4, 1)),
# 512 x 32 x 8
nn.Conv2d(512, 1024, kernel_size=f(1, 8), bias=False),
nn.BatchNorm2d(1024),
nn.ReLU(True),
# 1024 x 32 x 1
)
self.up_convs = nn.Sequential(
# 1024 x 32 x 1
nn.ConvTranspose2d(1024, 512, kernel_size=f(1, 8), bias=False),
nn.BatchNorm2d(512),
nn.ReLU(True),
# 512 x 32 x 8
self._UpBlock(512, 512, dilation=f(4, 1)),
self._UpBlock(512,
256,
stride=f(1, 2),
dilation=f(2, 1),
dilation_in=f(4, 1),
output_padding=f(0, 1)),
# 256 x 32 x 16
self._UpBlock(256, 256, dilation=f(2, 1)),
self._UpBlock(256,
128,
stride=f(1, 2),
dilation=f(1, 1),
dilation_in=f(2, 1),
output_padding=f(0, 1)),
# 128 x 32 x 32
nn.Conv2d(128, n_joints, kernel_size=1, bias=False),
# n x 32 x 32
)
init_parameters(self)
def __init__(self, n_joints, heatmap_space):
super().__init__()
self.n_joints = n_joints
self.heatmap_space = heatmap_space
self.down_layers = nn.Sequential(
self._regular_block(128, 128),
self._regular_block(128, 128),
self._down_stride_block(128, 192),
self._regular_block(192, 192),
self._regular_block(192, 192),
)
self.up_layers = nn.Sequential(
self._regular_block(192, 192),
self._regular_block(192, 192),
self._up_stride_block(192, 128),
self._regular_block(128, 128),
self._regular_block(128, self.n_joints),
)
init_parameters(self)
def __init__(self, resnet, n_joints):
super().__init__()
layers = [resnet.layer3, resnet.layer4]
for i, layer in enumerate(layers):
dilx = dily = 2**(i + 1)
for module in layer.modules():
if isinstance(module, nn.Conv2d):
if module.stride == (2, 2):
module.stride = (1, 1)
elif module.kernel_size == (3, 3):
kx, ky = module.kernel_size
module.dilation = (dilx, dily)
module.padding = ((dilx * (kx - 1) + 1) // 2,
(dily * (ky - 1) + 1) // 2)
self.layer1, self.layer2 = layers
self.hm_conv = nn.Conv2d(512, n_joints, kernel_size=1, bias=False)
init_parameters(self.hm_conv)
def __init__(self, n_joints):
super().__init__()
self.conv = nn.Conv2d(n_joints * 3, 128, kernel_size=1, bias=False)
init_parameters(self)
