向量化代码并从 pytorch 代码中删除嵌套循环答案

【问题标题】：Vectorize code and remove nested loops from pytorch code向量化代码并从 pytorch 代码中删除嵌套循环
【发布时间】：2021-08-13 15:21:02
【问题描述】：

是否可以删除嵌套循环以迭代所有图像，然后逐行迭代这些图像的所有像素？我发现这很困难，因为我需要使用另一个变量标签中像素（h，w）的位置来准确选择哪个通道的 h，w 值需要归零。

总的来说，我需要做的是给出分割输出（[2,2,2,2]）（n，通道，高度，宽度）和相应的标签（[2,2,2]）（n，高度宽度）。我需要添加另一个与输出相同大小的张量 M，除了所有像素的地面实况通道（来自标签）外，它在所有通道上都有一个恒定值 m。

【问题讨论】：

【解决方案1】：

您可以将其矢量化。认识到n、h 和w 对应于m 和labels 的所有像素值。因此，您可以执行以下操作：

m[:, labels, ...] = 0

【讨论】：

【解决方案2】：

知道我找到的答案是使用 torch.Tensor.scatter_()

new_labels = torch.unsqueeze(labels, 1)
m.scatter_(dim=1, index=new_labels, src=torch.zeros( 2, 2, 2, 2))

基本上，首先将标签的形状设置为 M 形状，然后通过选择轴在标签指定的位置为 m 添加零。

【讨论】：