同时指定一个numpy / torch轴的多个不同切片

Question

我有一批图像（尺寸为“ batchsize x通道x高度x宽度”的4d张量/阵列，我想在每个图像上绘制大小为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