维度不同的 Tensor 也能相加 ？

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

维度不同的 Tensor 如何相加呢？但是是自动填充对应维度，使维度相同后再相加。

如 torch.Size([1, 2, 3, 4]) 与 torch.Size([3, 4]) 相加，先将后者填充为 torch.Size([1, 2, 3, 4]) ，二者相加后 shape 为 [1, 2, 3, 4]

torch.Size([3, 2, 1, 4]) 与 torch.Size([3, 4]) 相加，则先将两者填充为 torch.Size([3, 2, 3, 4]) ，二者相加后 shape 为 [3, 2, 3, 4]，有种类似取最小公倍数的感觉。

代码示例1

import torch

a=torch.arange(1*2*3*4).reshape([1,2,3,4])    # torch.Size([1, 2, 3, 4])
b=torch.arange(3*4).reshape([3,4])      # torch.Size([3, 4])

print(f'a: {a}')
print(f'b: {b}')
print(f'(a+b): {a+b}')
print(f'(a+b).shape: {(a+b).shape}')     # torch.Size([1, 2, 3, 4])

输出结果：

a: tensor([[[[ 0,  1,  2,  3],
          [ 4,  5,  6,  7],
          [ 8,  9, 10, 11]],

         [[12, 13, 14, 15],
          [16, 17, 18, 19],
          [20, 21, 22, 23]]]])
          
b: tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
        
(a+b): tensor([[[[ 0,  2,  4,  6],
          [ 8, 10, 12, 14],
          [16, 18, 20, 22]],

         [[12, 14, 16, 18],
          [20, 22, 24, 26],
          [28, 30, 32, 34]]]])
(a+b).shape: torch.Size([1, 2, 3, 4])

代码示例2

a=torch.arange(1*2*3*4).reshape([3,2,1,4])    # torch.Size([3, 2, 1, 4])
b=torch.arange(3*4).reshape([3,4])      # torch.Size([3, 4])

print(f'a: {a}')
print(f'b: {b}')
print(f'(a+b): {a+b}')
print(f'(a+b).shape: {(a+b).shape}')     # torch.Size([3, 2, 3, 4])

输出结果：

a: tensor([[[[ 0,  1,  2,  3]],
         [[ 4,  5,  6,  7]]],

        [[[ 8,  9, 10, 11]],
         [[12, 13, 14, 15]]],

        [[[16, 17, 18, 19]],
         [[20, 21, 22, 23]]]])
         
b: tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
        
(a+b): tensor([[[[ 0,  2,  4,  6],
          [ 4,  6,  8, 10],
          [ 8, 10, 12, 14]],

         [[ 4,  6,  8, 10],
          [ 8, 10, 12, 14],
          [12, 14, 16, 18]]],


        [[[ 8, 10, 12, 14],
          [12, 14, 16, 18],
          [16, 18, 20, 22]],

         [[12, 14, 16, 18],
          [16, 18, 20, 22],
          [20, 22, 24, 26]]],


        [[[16, 18, 20, 22],
          [20, 22, 24, 26],
          [24, 26, 28, 30]],

         [[20, 22, 24, 26],
          [24, 26, 28, 30],
          [28, 30, 32, 34]]]])
(a+b).shape: torch.Size([3, 2, 3, 4])