def test_pixelshuffle2d():
nchan = 2
up_x = 2
up_y = 3
nx = 2
ny = 3
shape_before = (1, nchan * up_x * up_y, nx, ny)
shape_after = (1, nchan, nx * up_x, ny * up_y)
layer = PixelShuffle2D((up_x, up_y))
x = mx.nd.arange(np.prod(shape_before)).reshape(shape_before)
y = layer(x)
assert y.shape == shape_after
# - Channels are reshaped to form 2x3 blocks
# - Within each block, the increment is `nx * ny` when increasing the column
# index by 1
# - Increasing the block index adds an offset of 1
# - Increasing the channel index adds an offset of `nx * up_x * ny * up_y`
assert_allclose(
y,
[[[[ 0, 6, 12, 1, 7, 13, 2, 8, 14],
[18, 24, 30, 19, 25, 31, 20, 26, 32],
[ 3, 9, 15, 4, 10, 16, 5, 11, 17],
[21, 27, 33, 22, 28, 34, 23, 29, 35]],
[[36, 42, 48, 37, 43, 49, 38, 44, 50],
[54, 60, 66, 55, 61, 67, 56, 62, 68],
[39, 45, 51, 40, 46, 52, 41, 47, 53],
[57, 63, 69, 58, 64, 70, 59, 65, 71]]]]
)
def __init__(self,
in_channels,
out_channels,
scale_factor,
bn_use_global_stats,
bn_cudnn_off,
**kwargs):
super(DucBlock, self).__init__(**kwargs)
mid_channels = 4 * out_channels
with self.name_scope():
self.conv = conv3x3_block(
in_channels=in_channels,
out_channels=mid_channels,
bn_use_global_stats=bn_use_global_stats,
bn_cudnn_off=bn_cudnn_off)
self.pix_shuffle = PixelShuffle2D(factor=scale_factor)
def __init__(self,
backbone,
backbone_out_channels,
channels,
bn_use_global_stats=False,
bn_cudnn_off=False,
return_heatmap=False,
fixed_size=True,
in_channels=3,
in_size=(256, 192),
keypoints=17,
**kwargs):
super(AlphaPose, self).__init__(**kwargs)
assert (in_channels == 3)
self.in_size = in_size
self.keypoints = keypoints
self.return_heatmap = return_heatmap
with self.name_scope():
self.backbone = backbone
self.decoder = nn.HybridSequential(prefix="")
self.decoder.add(PixelShuffle2D(factor=2))
in_channels = backbone_out_channels // 4
for out_channels in channels:
self.decoder.add(
DucBlock(in_channels=in_channels,
out_channels=out_channels,
scale_factor=2,
bn_use_global_stats=bn_use_global_stats,
bn_cudnn_off=bn_cudnn_off))
in_channels = out_channels
self.decoder.add(
conv3x3(in_channels=in_channels,
out_channels=keypoints,
use_bias=True))
self.heatmap_max_det = HeatmapMaxDetBlock(
channels=keypoints,
in_size=(in_size[0] // 4, in_size[1] // 4),
fixed_size=fixed_size)