PyTorch 中张量运算广播

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]

《PyTorch 官方文档:BROADCASTING SEMANTICS》

相关推荐
AI技术学长3 分钟前
训练神经网络的批量标准化(使用 PyTorch)
人工智能·pytorch·神经网络·数据科学·计算机技术·批量标准化
ccLianLian4 分钟前
深度学习·经典模型·Transformer
人工智能·深度学习·transformer
belldeep9 分钟前
python:sklearn 决策树(Decision Tree)
python·决策树·机器学习·sklearn
Chat_zhanggong34516 分钟前
deepseek海思SD3403边缘计算AI产品系统
人工智能·边缘计算
人大博士的交易之路26 分钟前
龙虎榜——20250429
人工智能·数学建模·数据挖掘·量化分析·涨停回马枪·龙虎榜
哲讯智能科技40 分钟前
无锡哲讯科技:SAP财务系统——赋能企业智慧财务管理
大数据·人工智能
Want5951 小时前
DeepSeek: 探索未来的深度学习搜索引擎
人工智能·深度学习·搜索引擎·aigc
Tomorrow'sThinker1 小时前
基于PyTorch的图像分类特征提取与模型训练文档
人工智能·pytorch·分类
Kairo_011 小时前
A2A与MCP:理解它们的区别以及何时使用
python
wei3872452321 小时前
java练习2
java·开发语言·python