pytorch 2.0初探：和pytorch 1.13的速度对比

看到pytorch2.0出来了，而且宣传提速明显，一行代码即可提速43%左右：

compiled_model = torch.compile(model)

We then measure speedups and validate accuracy across these models. Since speedups can be dependent on data-type, we measure speedups on both float32 and Automatic Mixed Precision (AMP). We report an uneven weighted average speedup of 0.75 * AMP + 0.25 * float32 since we find AMP is more common in practice.

Across these 163 open-source models torch.compile works 93% of time, and the model runs 43% faster in training on an NVIDIA A100 GPU. At Float32 precision, it runs 21% faster on average and at AMP Precision it runs 51% faster on average.

因为想要提高炼丹速度，所以尝试和pytorch1.13进行对比。

环境

通过nvidia-smi查看显卡和驱动：

 显卡是4070ti，驱动安装的是最新的（好像是2023.6月的）。系统是Ubuntu 20.04。

conda环境有两个，分别是pytorch1.13和pytorch2.01，均是cuda 11.7版本，python是3.8.17。

 直接按照官方安装教程就行，这里我用的是pip进行安装的。

训练速度

代码用的是自己训练车牌识别的项目，主干网络用的是resnet34，代码中使用了AMP。

 pytorch 1.13

可以看到训练一轮耗时246s，中间过程太长没有全截下来。

准确点应该训练多轮取平均时间，我大致看了几轮，耗时基本一致，就没浪费时间跑太多轮。

pytorch 2.0.1

先不使用torch.compile

 

 耗时227s，相比pytorch1.13，提速8%左右，还是挺惊喜的，毕竟还没用torch.compile，要是用了不得上天啊，想想还挺激动的。

下面使用torch.compile，方法很简单，使用最上面那行代码就行，开搞！

耗时226s，不能说一模一样，简直就是毫无差别。 

然后注意到刚开始训练时，有一个警告：

anaconda3/envs/pt2/lib/python3.8/site-packages/torch/_inductor/compile_fx.py:90: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn( 

torch.set_float32_matmul_precision('high')这个接口有三个参数，默认是highest，可以通过torch.get_float32_matmul_precision()查看，于是我把这个参数设置成high，然后重新训练了，结果发现耗时还是一样。

然后去网上搜索一番，看到torch.compile有一个mode参数

• mode (str) – Can be either “default”, “reduce-overhead” or “max-autotune”

有人说reduce-overhead对小模型有效，于是我尝试把mode设置成reduce-overhead。

再次开搞。

 耗时227s，一点儿没变！！！

细心的我发现，刚开始训练时，又有个告警：

torch._inductor.utils: [WARNING] skipping cudagraphs due to input mutation 

这次不太幸运，Google了一圈也没发现什么原因，而且我还尝试了把mode参数设置成max-autotune，也没效果，而且多了一个警告：

先不折腾了，待后续观察。

总结

通过多次实验对比pytorch1.13和pytorch2.0.1，发现2.0.1比1.13提速8%左右，但是compile没有效果，暂时没搞清楚什么原因。

官方描述中有下面一段话：

Caveats: On a desktop-class GPU such as a NVIDIA 3090, we’ve measured that speedups are lower than on server-class GPUs such as A100. As of today, our default backend TorchInductor supports CPUs and NVIDIA Volta and Ampere GPUs. It does not (yet) support other GPUs, xPUs or older NVIDIA GPUs.

难道是40系显卡还不支持？

另外从官方介绍也可以看出，提速效果跟硬件、模型结构都有关系，所以具体能提速多少，还得自己在实际项目中尝试。