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

Question

我有一个形式为 (B, 2, H, W) 的火炬张量 pred，我想对两个不同的值 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.]]]])

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

How to add two torch tensors along given dimension?

python

broadcasting

torch