Torch中的广播机制

在Torch中，广播机制（Broadcasting）是一种强大的功能，它允许不同形状的张量在进行算术运算时自动扩展以匹配彼此的形状。这使得我们可以在不显式地复制数据的情况下进行操作，从而提高了计算效率和代码的简洁性。

我们可以使用非常简单的规则去描述这种广播机制：

1. 广播前提
   • 两个张量的从右向左每一对维度必须满足以下条件之一：
     • 维度相等
     • 其中一个维度为1
   • 当然，对于不同维度的张量，较小的张量会在左侧补充1，直到两个张量的维度数相等。
2. 广播规则
   • 在满足广播条件的前提下，从右向左，依次比较每个维度：
     • 如果两个维度相等，则保持不变。
     • 如果其中一个维度为1，则将其扩展为另一个张量的维度。

在扩展的时候，会将该维度的那个元素复制多次，直到与另一个张量的维度相等。
如

a = torch.tensor([1])  # shape: (1)
# 扩展到(3)
a = torch.tensor([1, 1, 1])  # shape: (3)
# 扩展到(2,3)
a = torch.tensor([[1, 1, 1], # 这是上面的(3)
                  [1, 1, 1]])  # shape: (2, 3)
# 扩展到(2,2,3)
a = torch.tensor([[[1, 1, 1],
                   [1, 1, 1]], # 这是上面的(2,3)
                  [[1, 1, 1],
                   [1, 1, 1]]])  # shape: (2, 2, 3)

下面我们给出几个例子：

加法运算

两个张量有一个维度相同，仅拓展一个张量

import torch

a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])  # shape: (2, 3)

b = torch.tensor([10, 20, 30])  # shape: (3)

# 对齐形状：
# a: (2, 3)
# b: (1, 3)，左侧补 1
# 结果：(2, 3)
# 扩展的 b:
# [[10, 20, 30],
#  [10, 20, 30]]

print(a + b)
# tensor([[11, 22, 33],
#         [14, 25, 36]])

两个张量都要拓展

a = torch.tensor([[1],
                  [2]])         # shape: (2, 1)

b = torch.tensor([[10, 20, 30]])  # shape: (1, 3)

# 对齐形状：
# a: (2, 1) -> (2, 3)
# b: (1, 3) -> (2, 3)
# 结果：(2, 3)
# a: [[1],[2]] -> [[1, 1, 1],
#                 [2, 2, 2]] (2, 3)
# b: [[10, 20, 30]] -> [[10, 20, 30],
#                      [10, 20, 30]] (2, 3)

print(a + b)
# tensor([[11, 21, 31],
#         [12, 22, 32]])

张量与标量运算

a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])  # shape: (2, 3)

b = torch.tensor(10)           # shape: (), 零维张量，即标量

# b: ()，由上面的规则，没有的维度会不断补充1，即()-> (1, 1)
# 对齐形状：
# a: (2, 3)
# b: (1, 1) -> (2, 3)
# b: [[10, 10, 10],
#     [10, 10, 10]] (2, 3)

# 标量可以参与每一个元素的运算。
print(a + b)
# tensor([[11, 12, 13],
#         [14, 15, 16]])

乘积

在Torch中，对于两个张量，使用如下表达式

a = torch.tensor([[1, 2, 3],
                  [1, 2, 3],
                  [1, 2, 3]])  # shape: (3, 3)
b = torch.tensor([[2, 2, 2],
                  [2, 2, 2],
                  [2, 2, 2]])  # shape: (3, 3)
ans = a * b

# ans = [[2, 4, 6],
#        [2, 4, 6],
#        [2, 4, 6]]

上述的a * b并不是传统意义上的矩阵乘法，而是被称为Hadamard积，即对应元素相乘。对于两个形状相同的张量，Hadamard积就是对应元素相乘。而在torch中，矩阵乘法应使用A @ B或者torch.matmul(A, B)来实现。

下面是一个使用广播机制的hadamard积的例子：

import torch

a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])  # shape: (2, 3)

b = torch.tensor([10, 20, 30])  # shape: (3)

result = a * b

print(result)
# tensor([[ 10,  40,  90],
#         [ 40, 100, 180]])