PyTorch：使用 1-D 张量和多通道 (3-D) 图像张量执行 add/sub/mul 操作

Question

注意：我正在寻找最快/最佳的方式来做或改进这一点，因为我的限制是时间

使用 PyTorch，我将图像作为 3D 张量，假设尺寸为 64 x 400 x 400，其中 64 是指通道，400 是图像尺寸。有了这个，我有一个 64 长度的 1D 张量，都有不同的值，目的是每个通道使用一个值。我想使用每个 1D 张量的值来应用于通道的整个 400x400 块。例如，当我想添加3d_tensor + 1d_tensor 时，我希望将1d_tensor[i] 添加到3d_tensor[i] 中的所有400x400 = 160000 值中，[i] 的范围从0 到63。

我之前做过的事情：
我试过直接做，只使用运算符：

output_add = 1d_tensor + 3d_tensor

这返回了3d_tensor(400)和1d_tensor(64)的维度不兼容的错误。

所以我目前的形式是使用 for 循环

for a, b in zip(3d_tensor, 1d_tensor):
   a += b

但是在一个阶段，我有四个不同的一维张量可以同时用于加法、减法或乘法，那么这种 for 循环方法是最有效的吗？我还计划每张图像执行 20 多次，所以速度是关键。我尝试将一维张量也扩展到 64 x 400 x 400 的尺寸，所以它可以直接使用，但我无法使用tensor.repeat()

Answer 1

您应该为一维数组添加一些维度：从 64 转换为 (64 x 1 x 1)：

output_add = 1d_tensor[:, None, None] + 3d_tensor

使用此无类型索引，您可以在任何位置添加维度。 [:, None, None] 将在现有维度之后向一维数组添加两个额外的维度。

或者您可以使用view 获得相同的结果，请参阅：

output_add = 1d_tensor.view(-1, 1, 1) + 3d_tensor

结果与 3D 数组具有相同的维度：(64, 1, 1)。所以pytorch可以使用广播。

这里是广播的一个很好的解释：How does pytorch broadcasting work?