py广播机制

广播机制

a = torch.arange(3).reshape(3, 1)
b = torch.arange(2).reshape(1, 2)
a+b

输出

tensor([[0, 1],
        [1, 2],
        [2, 3]])

将一个3 x 1的列向量与一个1 x 2的行向量进行相加，结果为3*2的矩阵。

广播的原则

• 扩展维度为1的张量。
• 低维度像高维度扩展

A = torch.arange(6).reshape(3, 2)
B = torch.arange(4).reshape(2,1, 2)
A,B,A+B
(tensor([[0, 1],
         [2, 3],
         [4, 5]]),
 tensor([[[0, 1]],
 
         [[2, 3]]]),
 tensor([[[0, 2],
          [2, 4],
          [4, 6]],
 
         [[2, 4],
          [4, 6],
          [6, 8]]]))