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

Posted 2023-03-29

【中文标题】同时分配 numpy/torch 轴的多个不同切片

【英文标题】：Assigning multiple different slices of a numpy/torch axis simultaneously

【发布时间】：2020-02-29 02:34:44

【问题描述】：

我有一批图像（尺寸为“batchsize x 通道 x 高度 x 宽度”的 4d 张量/数组，我想在每个图像上绘制大小为零的水平条，但在每个图像的不同行上绘制。我可以使用 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批次中两个图像的这些部分。我如何抓取这些不同的部分，但每张图像只有一个？在这种情况下，如果输入是 2xCxHxW，我想要 2xCx2xW，其中第一项对应于第一张图像的第 1 行和第 2 行，第二项对应于第二张图像的第 2 行和第 3 行。谢谢你。

【参考方案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