pytorch张量运算的广播机制

PyTorch 的广播机制(broadcasting)是指在进行张量运算时,自动扩展较小张量的形状以匹配较大张量的形状,使它们能够进行逐元素运算。广播机制避免了手动扩展张量的繁琐过程,并且在不增加内存开销的情况下进行高效计算。

广播规则

  1. 比较张量的形状:从后向前比较两个张量的每个维度(即从最右边的维度开始)。
  2. 维度匹配
    • 如果两个维度相等,则可以进行运算。
    • 如果一个张量在该维度上为 1,另一个张量为任意数值,则形状为 1 的张量会沿着该维度扩展,以匹配另一个张量的形状。
    • 如果两个张量在某个维度上不相等且没有一个是 1,则无法进行广播,运算会抛出错误。

广播机制的示例代码

1. 标量与张量的运算
复制代码
import torch

# 标量和张量相加
scalar = torch.tensor(3)
tensor = torch.tensor([1, 2, 3])

result = scalar + tensor
print(result)  # 输出: tensor([4, 5, 6])
2. 不同形状的张量运算
复制代码
import torch

# 创建两个形状不同的张量
A = torch.tensor([[1, 2, 3], [4, 5, 6]])  # 形状为 [2, 3]
B = torch.tensor([1, 2, 3])              # 形状为 [3]

# B 张量的形状会沿着第一个维度自动扩展为 [2, 3]
result = A + B 
print(result)   # 输出: tensor([[2, 4, 6], [5, 7, 9]])
3. 高维张量的广播
复制代码
import torch

# 形状为 [2, 1, 3]
C = torch.tensor([[[1, 2, 3]], [[4, 5, 6]]])

# 形状为 [3]
D = torch.tensor([1, 2, 3])

# D 张量的形状会沿着第一个和第二个维度扩展为 [2, 1, 3]
result = C + D  
print(result)   # 输出: tensor([[[ 2,  4,  6]], [[ 5,  7,  9]]])
4. 不兼容的张量运算
复制代码
import torch

# 形状为 [2, 3]
E = torch.tensor([[1, 2, 3], [4, 5, 6]])

# 形状为 [2]
F = torch.tensor([1, 2])

# 尝试进行运算将抛出错误,因为 E 和 F 在最后一个维度上不匹配
try:
    result = E + F  # 形状不兼容,会抛出错误
except RuntimeError as e:
    print("Error:", e)

总结

广播机制极大简化了张量运算的代码编写,特别是在处理不同行数和列数的张量时。理解广播机制能够帮助你编写更高效、简洁的代码,并充分利用 PyTorch 的计算能力。

相关推荐
蓝婷儿43 分钟前
Python 机器学习核心入门与实战进阶 Day 1 - 分类 vs 回归
python·机器学习·分类
Devil枫2 小时前
Kotlin扩展函数与属性
开发语言·python·kotlin
程序员阿超的博客3 小时前
Python 数据分析与机器学习入门 (八):用 Scikit-Learn 跑通第一个机器学习模型
python·机器学习·数据分析·scikit-learn·入门教程·python教程
xingshanchang4 小时前
PyTorch 不支持旧GPU的异常状态与解决方案:CUDNN_STATUS_NOT_SUPPORTED_ARCH_MISMATCH
人工智能·pytorch·python
reddingtons5 小时前
Adobe Firefly AI驱动设计:实用技巧与创新思维路径
大数据·人工智能·adobe·illustrator·photoshop·premiere·indesign
CertiK5 小时前
IBW 2025: CertiK首席商务官出席,探讨AI与Web3融合带来的安全挑战
人工智能·安全·web3
Deepoch6 小时前
Deepoc 大模型在无人机行业应用效果的方法
人工智能·科技·ai·语言模型·无人机
Deepoch6 小时前
Deepoc 大模型:无人机行业的智能变革引擎
人工智能·科技·算法·ai·动态规划·无人机
kngines6 小时前
【字节跳动】数据挖掘面试题0003:有一个文件,每一行是一个数字,如何用 MapReduce 进行排序和求每个用户每个页面停留时间
人工智能·数据挖掘·mapreduce·面试题
Binary_ey6 小时前
AR衍射光波导设计遇瓶颈,OAS 光学软件来破局
人工智能·软件需求·光学软件·光波导