如何沿给定维度添加两个火炬张量？

Question

我有一个火炬张量pred，格式为 (B, 2, H, W)，我想将两个不同的值 val1 和 val2 与轴 1 上的通道相加。

我设法通过直接访问单个频道以“机械”方式做到这一点，例如：

def thresh_format(pred, val1, val2):
    tr = torch.zeros_like(pred)
    tr[:, 0, :, :] = tr[:, 0, :, :].add(val1)
    tr[:, 1, :, :] = tr[:, 1, :, :].add(val2)
    return pred + tr

但是我想知道是否有“更好”的方法来做到这一点，例如通过利用广播。我对documentation 的理解是广播是从尾随维度发生的，所以在这种情况下，我很困惑如何使它适用于维度 1。 有什么想法吗？

Answer 1

实现这一点的最简单方法是将 val1 和 val2 堆叠在一个张量中，并对其进行整形以匹配 pred 张量沿公共维度的形状。

 pred + torch.tensor([val1, val2]).reshape((1,-1,1,1))

这样，对于加法，torch 会自动沿 pred 具有更高阶的维度广播值。

这与您将简单的标量值添加到张量时发生的情况非常相似，例如：

>>> torch.ones((2, 2)) + 3.
tensor([[4., 4.],
        [4., 4.]])

但不是在加法过程中向张量的每个元素广播一个标量值，在上述情况下，这些值是沿着尚未匹配的维度广播的。

>>> B=1; W=2; H=2; val1=3; val2=7
>>> pred = torch.zeros((B,2,W,H))
>>> val = torch.tensor([val1, val2]).reshape((1,-1,1,1))
>>> pred
tensor([[[[0., 0.],
          [0., 0.]],

         [[0., 0.],
          [0., 0.]]]])
>>> val
tensor([[[[3]],

         [[7]]]])
>>> pred + val
tensor([[[[3., 3.],
          [3., 3.]],

         [[7., 7.],
          [7., 7.]]]])