电脑基础 · 2023年3月31日

Pytorch 2.0来了!来看看新特性怎么应用到自己的代码里

文章目录

    • 1. 前言
    • 2. 安装和使用Pytorch 2.0
    • 3. 结语

1. 前言

Pytorch2.0和GPT4、文心一言同一时间段发布,可谓是热闹至极,我看了看Pytorch 2.0的文档,一句话概括下,2.0的功能介绍,核心就是torch.compile

opt_module = torch.Compile(module)

加入这行代码就能优化你的模型,优化后的模型和往常使用方式一样,推理速度会提升,比较重要的一点是,可以用于训练或者部署,训练可以传梯度,这次是带有AOTautograd的。然而需要注意的是,这行代码(编译)本身会消耗不少时间。

Pytorch官方在A100上测试了三个模型仓库的模型,加速比如下:
Pytorch 2.0来了!来看看新特性怎么应用到自己的代码里

Speedups for torch.compile against eager mode on an NVIDIA A100 GPU

看起来很不错。那废话不多说,来看看怎么用。

2. 安装和使用Pytorch 2.0

要安装pytorch2.0首先得上pytorch官网:
Pytorch 2.0来了!来看看新特性怎么应用到自己的代码里
需要注意的是,Pytorch 2.0只支持python3.8以上的版本,老环境记得升级一下。

按照官方提供的渠道安装完torch之后,就可以应用到自己的代码里了。这里我简单写一个用来做图像分类的神经网络。

import torch
import torch.nn as nn
import torch._dynamo
# ------ 加入这行,查看编译步骤 ------ #
torch._dynamo.config.verbose=True
# ------ 加入这行,跳过异常 ------ #
torch._dynamo.config.suppress_errors = True
class ThdCNN(nn.Module):
	# 网络结构
def fit(**args):
	# 训练和测试
model= ThdCNN(image_size, num_classes) # 声明和定义神经网络
# ------ 加入这行代码,model可选:default(compile), reduce-overhead, max-autotune ------ #
opt_model = torch.compile(model, mode="default")
optimizer = torch.optim.Adam(opt_model.parameters()) # 定义优化器
fit(opt_model, **args) # 训练和测试

我对三种加速模式进行了一个Epoch的测试(编译用时+训练和测试用时),结果如下:

模型 无加速 default(compile) reduce-overhead max-autotune
CNN(CPU) 101s 8s+78s 8s+85s 9s+81s
CNN(CUDA) 92s 7s+77s 8s+75s 8s+71s

可能是由于模型太小了,所以三种加速模式的效果差距不大,但按照官方的描述,三种模式的适用情况应该是:

  1. default(compile):针对大模型优化,编译时间短,没有额外的内存使用;
  2. reduce-overhead:针对小模型优化,用于减少框架开销,会使用一些额外的内存;
  3. max-autotune:模型整体优化,用于生成最优模型,但编译时间长。

3. 结语

Pytorch2.0与1.x是完全兼容的,因此基本上不会出现迁移错误,不过新特性还有很多没有玩明白的特性,所以在写代码的时候如果出现其他问题,可以参照 故障排除 解决。

compile的控制参数:

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  # advanced backend options go here as kwargs
  **kwargs
) -> torch._dynamo.NNOptimizedModule
  1. mode 指定编译器在编译时应该优化什么,包括efault(compile), reduce-overhead, max-autotune。
  2. dynamic 是否启用动态形状追踪。
  3. fullgraph 类似于 Numba 的nopython. 是否将模型分解成几个子图。
  4. backend要使用的后端。默认情况下,使用 TorchInductor,所有的后端可以通过torchdynamo.list_backends()查看。