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》

相关推荐
小黄人20257 分钟前
自动驾驶安全技术的演进与NVIDIA的创新实践
人工智能·安全·自动驾驶
ZStack开发者社区1 小时前
首批 | 云轴科技ZStack加入施耐德电气技术本地化创新生态
人工智能·科技·云计算
失败又激情的man2 小时前
python之requests库解析
开发语言·爬虫·python
打酱油的;2 小时前
爬虫-request处理get
爬虫·python·django
X Y O2 小时前
神经网络初步学习3——数据与损失
人工智能·神经网络·学习
唯创知音2 小时前
玩具语音方案选型决策OTP vs Flash 的成本功耗与灵活性
人工智能·语音识别
Jamence2 小时前
多模态大语言模型arxiv论文略读(151)
论文阅读·人工智能·语言模型·自然语言处理·论文笔记
tongxianchao2 小时前
LaCo: Large Language Model Pruning via Layer Collapse
人工智能·语言模型·剪枝
HyperAI超神经3 小时前
OmniGen2 多模态推理×自我纠正双引擎,引领图像生成新范式;95 万分类标签!TreeOfLife-200M 解锁物种认知新维度
人工智能·数据挖掘·数据集·图像生成·医疗健康·在线教程·数学代码
网安INF3 小时前
深度学习中批标准化与神经网络调优
人工智能·深度学习·神经网络·机器学习