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

Posted 2023-03-27

【中文标题】向量化代码并从 pytorch 代码中删除嵌套循环

【英文标题】：Vectorize code and remove nested loops from pytorch code

【发布时间】：2021-08-13 15:21:02

【问题描述】：

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

Colab Link

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

【参考方案1】：

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

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

【讨论】：

我试过用这个，但也许我的解释是错误的，它到底应该是什么 '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 添加零。