最近把之前写的一个图像分类项目从PyTorch 1.13迁移到2.0,听说torch.compile能白嫖加速,就试着给模型加了个@torch.compile。结果发现训练速度反而慢了20%,报错还一堆,比如什么“dynamic shape not supported”。我用的就是标准ResNet50,输入尺寸固定,数据加载也没啥花活。是不是我姿势不对?还是说这玩意儿只对某些特定场景有效?求有经验的老哥指点一下,到底该怎么用compile才能不翻车?
PyTorch 2.0的torch.compile到底能不能直接加速老项目?踩坑了
全部回复
共 152 条这题我熟,之前把检测模型直接套compile也是慢得离谱,后来发现是dataloader里num_workers太少,GPU吃不满反而在等数据。你可以先试试把torch.compile放到训练循环外面,只编译forward部分,别整个model一起包,另外batch size调大一倍看看,我这边ResNet50在A100上这么弄大概能快15%左右。
其实这玩意儿对静态shape和计算密集型网络收益最明显,如果你的loss里有些动态操作或者自定义op,很容易触发graph break,那还不如不编译。建议先用torch.compile的mode="reduce-overhead"跑几个epoch对比下,实在不行就只在eval阶段开启,训练还是老老实实用原版,等后续版本把dynamic shape支持完善了再上。
我踩坑时发现编译时间也算在第一次迭代里,如果你用timeit统计总时长,前几十步的编译开销会拉低平均速度。可以先跑个50步热身再开始计时,另外试试把torch._dynamo.config.supports_dynamic_shape设为True,虽然慢一点但能避免不少报错。
torch.compile更适合大模型和动态图,小项目直接上反而容易因编译开销拖慢速度,关掉试试。
说实话我一开始也跟你一样,直接无脑加@torch.compile然后等着起飞,结果被dynamic shape的报错按在地上摩擦。后来仔细看了下官方文档,才发现这玩意儿默认会做很多静态假设,比如你的tensor shape和stride在trace期间得完全一致,哪怕你输入尺寸固定,只要中间有个reshape或者slice操作,它都可能觉得是动态的。我建议你先把torch._dynamo的日志打开,看看它到底在哪个节点上放弃编译了,很多时候是数据加载那边返回的batch维度带了额外信息,或者你的自定义loss里有tensor.item()这种强制同步的操作。另外,老项目里如果用了大量Python控制流,比如if tensor.sum() > 0这种,也特别容易触发graph break,一break性能就崩。还有一点,torch.compile在A100这种新卡上收益明显,但如果你是2080Ti或者V100,可能编译开销比省下的算力还大,特别是batch size小的时候,那个inductor的codegen时间都快赶上一次forward了。我现在习惯是先跑一次warmup,把编译时间算进总时长里,然后对比50个epoch的实际墙钟时间,而不是看单step的loss下降速度。如果你ResNet50输入固定,理论上应该能加速的,我猜你可能是没给model.eval()或者没把torch.no_grad()包对,导致training模式下某些buffer被当成动态依赖了。总之别指望白嫖,得花半天调,调通了确实能快个20-30%,但前提是你愿意跟它的编译日志死磕。
torch.compile对静态图场景确实有收益,但你得先确保输入尺寸和数据类型完全一致,另外batch size别在训练中变。我遇到过类似情况,后来发现把mode="reduce-overhead"改成默认的"default"反而更稳,动态shape报错多半是模型里有个别层输出了不定长维度。还有,先别急着compile整个模型,可以试试只包住backbone或者用torch._dynamo的源码调试模式看下graph被fallback到哪了,很多老代码里的Python控制流或自定义autograd.Function容易触发回退,一回退开销比不compile还大。你ResNet50如果数据加载没瓶颈,可能瓶颈在CPU预处理上,compile加速的是GPU计算部分,整体没提升不奇怪。
讲真,torch.compile这东西真不是无脑加的,我自己的体验是它对于小batch、固定shape的简单CNN反而容易负优化。你那个ResNet50训练变慢,大概率是编译开销摊不平,尤其如果epoch短或者step多但每步计算量不大,那点graph优化根本抵不上trace和codegen的时间。另外报dynamic shape那个错,八成是模型里某个op的输入维度被当成了运行时变量,比如flatten或者view里用了tensor.shape[0]这种,虽然你觉得自己输入固定,但中间层某些操作确实可能产生动态行为。我建议你先试试把torch.compile的mode调成reduce-overhead或者max-autotune,再不行就加参数dynamic=False,但更实际的是先跑个benchmark对比一下,如果收益小于5%就直接放弃,老项目真没必要为了这个折腾。还有个坑是编译后梯度检查会变怪,有些自定义loss或者hook会静默失效,得手动验证数值对不对。反正我的结论是,torch.compile更适合那种大模型、长序列、多卡并行或者推理场景,小项目迁移成本高收益低。
torch.compile这玩意儿真不是无脑加的,它默认会做很多图优化,但你那个报错多半是模型里有动态shape的隐藏操作,比如某些op的输出维度是运行时才确定的。我试过在分割模型上踩过类似的坑,后来把input的shape用静态tensor固定住,再把dynamo的dynamic参数显式设成False才跑通。速度提升的话,个人感觉ResNet这种CNN收益确实不大,反而transformer或者大模型效果明显,可能是计算密集度不够。你要是实在想用,可以先torch.compile(model, mode="reduce-overhead")试试,或者干脆开个profile看看瓶颈在哪,别一开始就指望白嫖。
这问题我太有同感了,之前也信了“无脑加速”的邪,结果小batch下compile反而更慢。你这种情况大概率是默认模式开了动态shape检测,加上图编译的预热开销在小模型上被放大了。建议试试torch.compile(model, mode="reduce-overhead"),或者把batch size调大点,收益会明显一些。另外老项目里如果有自定义loss或者数据增强里用了tensor操作,很容易触发graph break,那基本就白编译了,可以看看编译后的warnings。
这情况我也遇到过,compile对固定shape收益不大,反而图优化开销拖后腿,试试关掉dynamic和把batch调大点。
先用torch.compile默认模式跑通再说,别一上来就加花活,老项目里很多自定义层和动态op都容易踩雷。
小模型和训练场景真没必要硬上compile,推理阶段用可能收益更大,先关掉动态shape试试。
建议先把torch._dynamo的日志打开看下graph break,ResNet50大概率是某些算子没走CUDA graph才变慢。
小batch或GPU利用率不高时compile收益很小,先把batch调大或看看GPU占用再说。
我之前也遇到过一模一样的情况,ResNet50固定输入,加了compile反而慢了。后来发现大概率是CUDA graph和算子融合在第一次迭代时开销太大,而且你如果开了grad mode,torch.compile默认会做一些重编译检查,老代码里某些写法比如在forward里动态改层属性就会触发重新编译,直接拖垮速度。
我后来试了个偏方,先把模型切成推理模式跑通,再把训练循环里的loss.backward()和optimizer.step()也包进compile区域,但只对conv和BN层做全图捕获,其他自定义模块保持原样,稍微好点。不过说实话,这玩意对图像分类这种计算密集型模型收益真的有限,我测下来最多也就5%提升,还得花一堆时间调配置。
你要是真想用,建议把torch.compile的mode设成max-autotune,然后加个torch._dynamo.config.suppress_errors=True先跑通再说。另外检查下你是不是用了DataLoader的pin_memory和non_blocking,这俩和compile的缓存机制偶尔会冲突,关掉试试。
最稳的路子其实是先跑一遍official的benchmark脚本,看看你的GPU架构是不是被支持的老卡,有些卡上Inductor后端生成的kernel反而比手写慢。如果实在不行,就退回1.13吧,老项目稳定压倒一切,没必要为了赶时髦折腾一晚上。
小模型加compile收益不大,你试试把torch._dynamo.mark_dynamic去掉,再把batch size调大点,2.0对动态shape优化很保守。
torch.compile这个事儿我也踩过类似的坑,后来发现它默认会做不少动态shape的假设,你固定尺寸反而得显式设一下dynamic=False或者把输入tuple化。另外训练慢20%可能是graph break太多,你试试把数据预处理和loss计算也包进compile的scope里,或者换mode="reduce-overhead"。老项目直接无脑加确实容易翻车,建议先拿profile看看编译时间占比,如果epoch太短的话编译开销根本摊不回来。
我之前也踩过这坑,小batch下compile反而更慢,建议你试试把batch调大点,或者用mode="max-autotune"看看。
torch.compile对老项目真不是无脑加个装饰器就行,你这情况大概率是图模式下的编译开销比执行收益还大,尤其小batch或者GPU没吃满的时候,慢20%很常见。我之前在固定尺寸的检测模型上试过,得配合dynamic=False显式关掉动态shape,还要把torch._dynamo.config.suppress_errors打开先看warning,不然一堆隐式tensor操作会触发回退到eager模式,反而更慢。建议你先用torch.compile(model, mode="reduce-overhead")这个档位,再把数据加载和loss计算都挪到编译图外面,如果还是慢就放弃,老实换回1.13的apex混合精度,收益更稳定。另外你训练循环里如果用了.item()或者numpy()这类同步操作,也会打断图优化,检查下有没有这类隐藏坑。
训练阶段确实容易负优化,试试只对推理部分compile,或者把dynamic=True加上,ResNet50固定shape不该报这个错的。
老项目直接套compile确实容易翻车,我试过几次也是这感觉。你那个慢20%大概率是编译开销没摊平,batch太小或者epoch太少,建议先把torch._dynamo的日志打开看看graph break都卡在哪。还有dynamic shape那个报错,虽然你输入尺寸固定,但可能loss或者dataloader里有些tensor维度在变,检查下有没有用list当batch或者padding。想快速验证加速效果的话,可以先用torch.compile(model, mode="reduce-overhead")跑个50步看看,别直接全量训练,另外记得设torch.set_float32_matmul_precision("high"),有时候光这个就能提不少速。
这情况太正常了,torch.compile不是无脑套上去就完事的。你那个dynamic shape的报错,多半是模型里有个别层输出了不定长的中间变量,哪怕你输入是固定的,某些op在trace的时候也会被判定成动态。我建议你先用torch.compile(model, mode="reduce-overhead")试试,这个模式对训练场景更友好,默认模式有时候反而会引入额外开销。另外,编译前的warmup很重要,你跑几个batch让CUDA graph和inductor缓存都热起来再计时,不然前几次迭代的编译时间会严重拖慢你看到的速度。还有个小坑,如果用了torch.cuda.amp或者GradScaler,最好把compile放到amp上下文外面,不然有些算子会被重复特化。老实说,像ResNet50这种结构规整的模型,在A100上可能也就快个10%左右,如果卡是消费级的,收益更小,甚至可能负优化。你可以先用torch.profiler对比一下编译前后的kernel耗时,如果发现大量时间花在triton的代码生成上,那说明图没融合好,不如把目光转向数据加载和loss计算那边,可能收益更大。
torch.compile对老项目真不是无脑加的,你那个慢20%太正常了,我试过几次发现它特别吃显存和CUDA图优化的余量,小batch下反而容易负优化。ResNet50这种静态图按理说应该友好,但你得先确认一下是不是开了mode="reduce-overhead"或者把dynamic=True参数显式关掉,有时候默认模式会保守地插入很多guard检查。另外报错的话,建议先用torch._dynamo.config.suppress_errors=True跑一遍看能不能fallback,能跑通再慢慢调,不然直接硬上肯定翻车。我自己的经验是,compile对那种带大量小算子、频繁调Python控制流的模型提升明显,纯卷积堆叠的反而没啥肉吃,你可以先用torch.compile的profile看看瓶颈到底在哪再决定要不要继续用。
先别急着怀疑姿势,这玩意儿对显存和计算图复杂度很敏感,你固定尺寸的ResNet50按理说是理想对象,但老项目里如果有任何自定义的loss或者数据增强里带了随机性,都可能让编译器保守化。我碰到的坑是它会把一些非标准操作重编译,每次迭代都触发graph break,速度直接崩。建议你先把torch.compile的mode设成"max-autotune"试试,同时把torch._dynamo.config.capture_scalar_outputs打开,很多隐藏的
torch.compile对老项目真不是无脑加一行就完事,它默认会做很多激进优化,但遇到dynamic shape或者某些自定义op就直接摆烂回退,反而拖慢速度。你ResNet50输入固定的话,试试在compile里加上mode="reduce-overhead"或者显式指定dynamic=False,另外确认一下你的数据加载和预处理有没有隐式改变张量形状的地方。我之前也是类似情况,后来把模型里一个自适应池化改成固定尺寸才正常,提速大概15%左右,但对比折腾的时间真不如直接换cudnn benchmark。还有个坑是batch size太小的话编译开销根本摊不平,建议先拿64以上的batch试。