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