如何用Torch.gather结合双索引数组实现批量图像无循环随机裁剪

需要从尺寸为5×10的图像中随机提取2×5的图像块，并且对批量内的所有图像执行这一随机提取操作。目前已有嵌套循环的朴素实现，但希望编写无嵌套循环的高效方案，但不清楚如何用torch.gather处理双索引数组（示例中的idx_h与idx_w）。

以下是朴素嵌套循环实现代码：

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)

高效实现方案 #

方法1：广播式直接索引（无需torch.gather） #

通过扩展索引维度实现广播匹配，直接完成批量索引，代码简洁高效：

import torch
b = 3 # batch size
h = 5 # height
w = 10 # width
crop_border = (3, 5) # 裁剪像素数（高，宽）
x = torch.arange(b * h * w).reshape(b, h, w)

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

# 生成每个batch对应的高、宽索引范围
idx_h = torch.stack([torch.arange(d, d + (h - crop_border[0])) for d in dh_]) # shape (b, 2)
idx_w = torch.stack([torch.arange(d, d + (w - crop_border[1])) for d in dw_]) # shape (b, 5)

# 扩展维度实现广播：idx_h -> (b, 2, 1)，idx_w -> (b, 1, 5)
idx_h_exp = idx_h.unsqueeze(-1)
idx_w_exp = idx_w.unsqueeze(1)

# 直接索引，自动匹配维度
cropped_x = x[torch.arange(b)[:, None, None], idx_h_exp, idx_w_exp]

print(cropped_x)

方法2：使用torch.gather分步处理双索引 #

如果需要用torch.gather，可以分两步分别处理高度和宽度维度的索引：

import torch
b = 3 # batch size
h = 5 # height
w = 10 # width
crop_border = (3, 5) # 裁剪像素数（高，宽）
x = torch.arange(b * h * w).reshape(b, h, w)

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

idx_h = torch.stack([torch.arange(d, d + (h - crop_border[0])) for d in dh_]) # shape (b, 2)
idx_w = torch.stack([torch.arange(d, d + (w - crop_border[1])) for d in dw_]) # shape (b, 5)

# 第一步：按高度维度gather，提取每个batch对应高度的所有列
x_h = torch.gather(x, dim=1, index=idx_h.unsqueeze(-1).repeat(1,1,w)) # shape (b, 2, w)

# 第二步：按宽度维度gather，提取对应列
cropped_x = torch.gather(x_h, dim=2, index=idx_w.unsqueeze(1).repeat(1,2,1)) # shape (b, 2, 5)

print(cropped_x)

说明 #

两种方法都彻底避免了嵌套循环，利用PyTorch的向量化操作提升计算效率，批量越大优势越明显。其中广播索引的写法更直观简洁，torch.gather的方式则更贴合对该API的使用需求。

内容的提问来源于stack exchange，提问作者iric