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)