TLDR
右对齐,空补一,从左往右依维运算
[m] + [x, y] = [m +x, m + y]
正文
以如下 a
b
两个 tensor 计算为例
python
a = torch.tensor([
[1],
[2],
[3],
])
b = torch.tensor([
[
[1, 2, 3],
],
[
[4, 5, 6],
],
[
[7, 8, 9],
],
])
# a.shape = (3, 1)
# b.shape = (3, 1, 3)
首先将两个 tensor 的 shape 右对齐
a( , 3, 1)
b(3, 1, 3)
判断两个 tensor 是否满足广播规则:
- tensor 至少有一个维度(比如
torch.tensor((0,))
便不符合本要求) - 检查上一步对齐的 tensor shape,要求两个 tensor 对应维度的大小:要么相同;要么其中一个为 1;要么其中一个为空
- 如果满足上述规则,则继续,否则报错
将对齐后空缺的维度设置为 1
a(1, 3, 1)
b(3, 1, 3)
其实就是对 a
进行了扩维,此时两个 tensor 为:
python
a = torch.tensor([
[
[1],
[2],
[3],
],
])
b = torch.tensor([
[
[1, 2, 3],
],
[
[4, 5, 6],
],
[
[7, 8, 9],
],
])
# a.shape = (1, 3, 1)
# b.shape = (3, 1, 3)
从左往右对两个 tensor 的每一个维度进行运算,按照以下规则:
- 如果大小相同,则直接进行运算即可(一一对应)
- 如果其中一个大小为 1,则使用这个元素与另一个 tensor 当前维度下的每个元素进行运算(本质是一个递归操作)
例如计算 a + b
(这两个 tensor 已经经过上述步骤处理,即维度已经相同)
python
# 1. 因为 a.shape[0] == 1,所以将 a[0] 分别与 b[0]、b[1]、b[2] 相加
[
a[0] + b[0],
a[0] + b[1],
a[0] + b[2],
]
# 2. 接下来继续往后计算,以 a[0] + b[0] 为例
# 因为 a[0].shape[0] = 3, b[0].shape[0] = 1,
# 所以将 b[0][0] 分别与 a[0][0]、a[0][1]、a[0][2] 相加
[
[ # a[0] + b[0]
a[0][0] + b[0][0],
a[0][1] + b[0][0],
a[0][2] + b[0][0],
],
[ # a[0] + b[1]
a[0][0] + b[1][0],
a[0][1] + b[1][0],
a[0][2] + b[1][0],
],
[ # a[0] + b[2]
a[0][0] + b[2][0],
a[0][1] + b[2][0],
a[0][2] + b[2][0],
],
]
# 3. 继续往后计算,以 a[0][0] + b[0][0] 为例
# 因为 a[0][0].shape[0] == 1,
# 所以将 a[0][0][0] 分别与 b[0][0][0]、b[0][0][1]、b[0][0][2] 相加
[
[ # a[0] + b[0]
[ # a[0][0] + b[0][0]
a[0][0][0] + b[0][0][0],
a[0][0][0] + b[0][0][1],
a[0][0][0] + b[0][0][2],
],
[ # a[0][1] + b[0][0]
a[0][1][0] + b[0][0][0],
a[0][1][0] + b[0][0][1],
a[0][1][0] + b[0][0][2],
],
[ # a[0][2] + b[0][0]
a[0][2][0] + b[0][0][0],
a[0][2][0] + b[0][0][1],
a[0][2][0] + b[0][0][2],
],
],
[ # a[0] + b[1]
[ # a[0][0] + b[1][0]
a[0][0][0] + b[1][0][0],
a[0][0][0] + b[1][0][1],
a[0][0][0] + b[1][0][2],
],
[ # a[0][1] + b[1][0]
a[0][1][0] + b[1][0][0],
a[0][1][0] + b[1][0][1],
a[0][1][0] + b[1][0][2],
],
[ # a[0][2] + b[1][0]
a[0][2][0] + b[1][0][0],
a[0][2][0] + b[1][0][1],
a[0][2][0] + b[1][0][2],
],
],
[ # a[0] + b[2]
[ # a[0][0] + b[2][0]
a[0][0][0] + b[2][0][0],
a[0][0][0] + b[2][0][1],
a[0][0][0] + b[2][0][2],
],
[ # a[0][1] + b[2][0]
a[0][1][0] + b[2][0][0],
a[0][1][0] + b[2][0][1],
a[0][1][0] + b[2][0][2],
],
[ # a[0][2] + b[2][0]
a[0][2][0] + b[2][0][0],
a[0][2][0] + b[2][0][1],
a[0][2][0] + b[2][0][2],
],
],
]
总结
右对齐 ,空补一 ,从左往右依维 (递归 )运算。
一个 tensor 的某个维度大小为 1 时的计算规则:[1] + [2, 3, 4] = [1 + 2, 1 + 3, 1 + 4]
。