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

Question

我有一批图像(尺寸为“批大小x通道x高度x宽”的4d张量/数组)，我希望在每幅图像上绘制大小为s的水平条形图，但每幅图像跨越不同的行。我可以使用for循环来实现这一点，但是我还没有找到向量化的实现。

理想情况下，我会生成一个一维张量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张量，因为它正在从-在批处理中的图像中-从我如何抓取这些不同的部分，但只有一个为每个图像？在这种情况下，如果输入是2xCxHxW，我需要2xCx2xW，其中第一个项对应于第一个映像的第1和第2行，第二个项对应于第二个映像的第2和第3行。谢谢。

Answer 1

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

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