最近把之前写的一个图像分类项目从PyTorch 1.13迁移到2.0,听说torch.compile能白嫖加速,就试着给模型加了个@torch.compile。结果发现训练速度反而慢了20%,报错还一堆,比如什么“dynamic shape not supported”。我用的就是标准ResNet50,输入尺寸固定,数据加载也没啥花活。是不是我姿势不对?还是说这玩意儿只对某些特定场景有效?求有经验的老哥指点一下,到底该怎么用compile才能不翻车?
楼主
23小时前
PyTorch 2.0的torch.compile到底能不能直接加速老项目?踩坑了
请 登录 后发表回复
全部回复
共 1 条
2楼
22分钟前
同踩过这个坑,torch.compile对静态图确实有加速,但ResNet50这种老模型在PyTorch 2.0上反而容易因为算子兼容性问题变慢,尤其是dynamic shape报错说明你的输入张量在某些维度上被检测成了可变长,可以试试torch._dynamo.config.dynamic_shapes=False强行关掉动态形状检测。另外编译模式用mode="reduce-overhead"或者先跑几个warmup步让缓存编译好,效果会好不少,但说实话小批量训练提升不明显,更推荐用在推理或大batch场景。