Instantly share code, notes, and snippets.

# musyoku/PixelShuffler.py Last active Mar 26, 2018

 import math import numpy as np batchsize = 2 r = 2 out_channels = 3 in_channels = r ** 2 * out_channels in_height = 3 in_width = 3 out_height = in_height * r out_width = in_width * r in_map = np.zeros((batchsize, in_channels, in_height, in_width), dtype=np.int32) for b in xrange(batchsize): for k in xrange(in_channels): for h in xrange(in_height): for w in xrange(in_width): in_map[b, k, h, w] = in_channels * in_height * in_width * b + in_height * in_width * k + in_height * h + w print in_map out_map = np.reshape(in_map, (batchsize, r, r, out_channels, in_height, in_width)) # print out_map out_map = np.transpose(out_map, (0, 3, 4, 1, 5, 2)) # print out_map out_map = np.reshape(out_map, (batchsize, out_channels, out_height, out_width)) print out_map # test for b in xrange(batchsize): for k in xrange(out_channels): for h in xrange(in_height * r): for w in xrange(in_width * r): _k = out_channels * r * (h % r) + out_channels * (w % r) + k _h = int(math.floor(h / float(r))) _w = int(math.floor(w / float(r))) assert out_map[b, k, h, w] == in_map[b, _k, _h, _w]