torch gather using two index arrays

Question

The goal is to extract a random 2x5 patch from a 5x10 image, and do so randomly for all images in a batch. Looking to write a faster implementation that avoids for loops. Haven't been able to figure out how to use the torch .gather operation with two index arrays (idx_h and idx_w in code example).

Naive for loop:

import torch
b = 3 # batch size
h = 5 # height
w = 10 # width
crop_border = (3, 5) # number of pixels (height, width) to crop 
x = torch.arange(b * h * w).reshape(b, h, w)
print(x)

dh_ = torch.randint(0, crop_border[0], size=(b,))
dw_ = torch.randint(0, crop_border[1], size=(b,))

_dh = h - (crop_border[0] - dh_)
_dw = w - (crop_border[1] - dw_)

idx_h = torch.stack([torch.arange(d_, _d) for d_, _d in zip(dh_, _dh)])
idx_w = torch.stack([torch.arange(d_, _d) for d_, _d in zip(dw_, _dw)])
print(idx_h, idx_w)

new_shape = (b, idx_h.shape[1], idx_w.shape[1])
cropped_x = torch.empty(new_shape)

for batch in range(b):
    for height in range(idx_h.shape[1]):
        for width in range(idx_w.shape[1]):
            cropped_x[batch, height, width] = x[
                batch, idx_h[batch, height], idx_w[batch, width]
            ]
print(cropped_x)

iric · Answer

Index arrays needed to be repeated and reshaped to work with gather operation. Fast_crop code based pytorch discussion: https://discuss.pytorch.org/t/similar-to-torch-gather-over-two-dimensions/118827

def fast_crop(x, idx1, idx2):
    """
    Compute
        x: N x B x V
        idx1: N x K matrix where idx1[i, j] is between [0, B)
        idx2: N x K matrix where idx2[i, j] is between [0, V)
    Return:
        cropped: N x K matrix where y[i, j] = x[i, idx1[i,j], idx2[i,j]]

    """
    x = x.contiguous()
    assert idx1.shape == idx2.shape
    lin_idx = idx2 + x.size(-1) * idx1
    x = x.view(-1, x.size(1) * x.size(2))
    lin_idx = lin_idx.view(-1, lin_idx.shape[1] * lin_idx.shape[2])

    cropped = x.gather(-1, lin_idx)
    return cropped.reshape(idx1.shape)


idx1 = torch.repeat_interleave(idx_h, idx_w.shape[1]).reshape(new_shape)
idx2 = torch.repeat_interleave(idx_w, idx_h.shape[1], dim=0).reshape(new_shape)
cropped = fast_crop(x, idx1, idx2)

(cropped == cropped_x).all()

Using realistic numbers for b = 100, h = 100, w = 130 and crop_border = (40, 95), a 10 trial run takes the for loop 32s while fast_crop only 0.043s.

torch gather using two index arrays

1 Answers