Say you have a gray image tensor of shape (1, 1, 128, 128)
. What I would like to do here is to sample in each h, w
dimension with stride=2
, which would then make 4 sub-images of size (1, 1, 64, 64)
depending on where the indexing starts. Then, I would like to batch them to finally form the tensor of size (4, 1, 64, 64)
. For now, I’m using double for loops which is presumably inefficient. Is there an efficient way to vectorize this using functions implemented in pytorch?
The function I’m using right now:
def img_subsample(img, ratio=2):
'''
ex) Subsamples img of size 128 * 128,
output is batched into 4 * 1 * 64 * 64
Args:
img: input of size ex) 128 * 128
ratio(int): Subsampling ratio
Returns: batched, subsampled img
'''
sz = img.shape[-1]
batched_output = torch.zeros([ratio**2, 1, (sz//ratio), (sz//ratio)])
for i in range(ratio):
for j in range(ratio):
batched_output[ratio*i+j, :, :, :] = img[i::ratio, j::ratio]
return batched_output