python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > PyTorch Eager与Compiled模式

PyTorch中的 Eager与Compiled模式详解

作者:@ysw

PyTorch中的EagerMode和CompiledMode是两种不同的执行模式,EagerMode即时执行,代码逻辑与Python一致,易用性高但运行时开销大;而CompiledMode通过编译优化,能在不牺牲灵活性的前提下大幅提升运行效率,本文介绍PyTorch中的Eager与Compiled模式,感兴趣的朋友一起看看吧

简单来说,PyTorch 中的 Eager 和 Compiled 是两种截然不同的执行模式:

1. Eager Mode(即时执行模式)

这是 PyTorch 最经典的运行方式(从 0.4 版本开始成为默认模式)。

核心机制:
代码即定义即运行。当你写下 c = a + b 时,Python 解释器会立刻调用 CUDA 或 CPU 的核函数执行加法,并返回结果张量。没有显式的“构建图”和“执行图”的分离阶段

优点:

缺点:

2. Compiled Mode(编译执行模式)

主要指 PyTorch 2.0 及以上版本推出的 torch.compile 技术。

核心机制:
编译器会“捕获”你的模型,将其转换为一个全局的静态计算图(使用 TorchDynamo 捕获,TorchInductor 后端生成代码),然后进行算子融合(如将相邻的 add 和 relu 合并为一个核函数)、消除冗余内存访问,最终生成针对 GPU(如使用 Triton 或 CUDA)或 CPU 高度优化的底层代码。

优点:

缺点:

3. 两者对比总结

维度Eager ModeCompiled Mode (torch.compile)
执行方式解释执行,遇操作即运行先捕获图,编译优化,再运行
速度基准速度通常更快(融合算子,减少开销)
调试非常方便(print, pdb较复杂(优化后的代码不易直接调试)
内存占用基准占用通常更低(算子融合减少中间变量)
启动时间无预热时间有一次性编译预热时间
适用场景开发调试、动态性强的小模型生产部署、训练大模型、固定结构的推理

4. 如何使用?

torch.compile 并非要取代 Eager Mode,而是提供了一个“免费午餐”式的优化手段。

你可以这样组合使用:

import torch
def my_model(x):
    # 复杂的 Python 逻辑(这部分在编译模式下也能处理)
    return x.relu() @ x.T
# 在开发阶段:用 eager 模式,方便打印和调试
out_dev = my_model(torch.randn(10, 10))
print(out_dev.shape)
# 在训练/部署阶段:用 compile 包裹一行即可加速
compiled_model = torch.compile(my_model)
out_prod = compiled_model(torch.randn(10, 10))

总结来说,Eager 是 PyTorch 易用性的基石,让你能像写普通 Python 代码一样写神经网络;而 Compiled 是在此基础上,通过编译技术将动态代码转化为高效的静态执行计划,在不牺牲太多灵活性的前提下大幅提升运行效率。

到此这篇关于PyTorch中的 Eager与Compiled模式详解的文章就介绍到这了,更多相关PyTorch Eager与Compiled模式内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家!

您可能感兴趣的文章:
阅读全文