同时分配 numpy/torch 轴的多个不同切片

Question

我有一批图像 (4d tensor/array，尺寸为 "batchsize x channels x height x width"，我想在每个图像上绘制大小为 s 的零水平条，但每个图像跨越不同的行。我可以用 for 循环简单地做到这一点，但我一直无法找出矢量化实现。

理想情况下，我会生成一个 "batchsize" 随机起点的一维张量 r，然后执行类似 t[:,:,r:r+s,:] = 0。如果我尝试这个，我会得到 TypeError: only integer scalar arrays can be converted to a scalar index

如果我做一个玩具示例，只是尝试从只有两张图像的批次中拉出两个不同的部分，做类似 t[:,:,torch.tensor(([1,2],[2,3])),:] 的事情，我会得到一个 5D 张量，因为它正在拉动 both 这些部分来自批处理中的 both 个图像。如何抓取那些不同的部分，但每张图片只有一个？在这种情况下，如果输入是 2xCxHxW，我会想要 2xCx2xW，其中第一项对应于第一张图像的第 1 行和第 2 行，第二项对应于第二张图像的第 2 行和第 3 行。谢谢你。

Answer 1

您可以使用此函数，它会创建一个掩码，您可以在其中根据索引在 y 或 x 轴上执行操作。您可以通过将索引的 x 值设置为它们的 y 索引来做到这一点。

bsg = sgs.data
device = sgs.device
bs, _, x, y = bsg.shape
max_y = y-size-1
rs = torch.randint(0, max_y, (bs,1), device=device)
m = torch.arange(y,device=device).repeat(bs, x)
gpumask = ((m < rs) | (m > (rs+size))).view(bs, 1, x, -1)
gpumask*bsg

同时分配 numpy/torch 轴的多个不同切片

Assigning multiple different slices of a numpy/torch axis simultaneously

python

gpu

numpy

torch