不盲目 | 客观、科学的看待Pytorch2.0更新的新特性，没你想的那么快！ - 文章 - 开发者社区

学 picture.image

最近两天都在推送 Pytorch2.0 的新闻，感慨下更新的真快呀，Pytorch1.13还没怎么捂热，这又来2.0了。

picture.image

虽然自己的文章还有 一堆坑待填 的，还是忍不住看了看官方关于Pytorch2.0的博客，同时也说说自己的看法吧。

首先一句话概括下，2.0的功能介绍[1]，核心就是torch.compile：

picture.image

一行代码就能优化你的模型，优化后的模型和往常使用方式一样，速度会提升，比较重要的一点是， 可以用于训练或者部署，训练可以传梯度， 这次是带有AOTautograd的 。

官方测试了三个模型仓库的模型，在A100的加速比如下：

picture.image

看提速，提升没有超一倍，相比TensorRT（TensorRT优化的话，一般都是2-6倍左右，当然要算FP16）这种变态来说，torch.compile优化级别还是差些的，性能肯定不是极致，更不如纯手工优化后的。

不过胜在好用以及拓展性强，可以应付大部分的训练和部署场景，使用起来很方便。而TensorRT这种的话，上手比较难些：内卷成啥了还不知道TensorRT？超详细入门指北，来看看吧！

不过值得吐槽的一点，现在和Pytorch沾边的库越来越多了，加速模型的，部署模型的，编译器类型的，我都认不全了。大概有这些：

torchscript(torch.jit)
lazy tensor
torch.fx
torch::deploy
AItemplate
nvfuser
TorchDynamo
TorchInductor
AOTAutograd
PrimTorch
torch_tensorrt
functorch
torch2trt

picture.image

而这次2.0强调的是 TorchDynamo 和 TorchInductor ，一个是解析Python字节码，可以trace你的model；一个是神经网络编译器，将pytorch模型lower为IR然后生成高性能的CUDA或者CPU代码。

TorchDynamo相比fx和jit的trace，可以更方便地跟踪data-dependent分支，而且不像torch.script一样需要修改代码，这个后续老潘会单独介绍下。

感觉TorchDynamo就是加强普适版的fx和torchscript，而TorchInductor就是编译器，有对应的IR（define-by-run IR），然后可以利用triton去codegen生成高性能的算子。

triton这个名字有点歧义哈，这里指的不是triton-server-inference，而是一个类似于TVMscript的可以通过python语法去写高性能GPU程序的，大家不要混了：

picture.image

无奈感慨下，深度学习编译器大繁荣的时代来了，啥都要编译来干了，不管是之前的torchscript还是torch.fx，以及新出的 TorchDynamo 和 TorchInductor， 总之就是编译优化编译优化：

picture.image

编译器这块晦涩难懂，想要了解底层学习的，难度还是挺大（没时间学了）。希望Pytorch这个新的编译器，之后使用起来bug少一些，问题看起来直观一些...

简单测试下

这个Pytorch2.0版本其实就是1.14，只不过官方认为torch.compile这个功能意义重大，所以干脆2.0了，现在（12月2日）放出的是beta版，稳定版的话在明年3月。所以说对之前版本完全兼容（TF2.0情何以堪）：

picture.image

首先按照官网的要求来安装：


      
        
pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

我自己的环境是cuda11.8版本，一开始测试的时候也有问题，以为是CUDA版本不匹配，单独测试了官方的镜像，发现也运行不起来（这个确实是坑，按照官方安装镜像直接运动提示缺很多东西）。后来还是在host下，看报错发现少了几个python库（比如tqdm），安装了就能跑了：


      
        
oldpan@oldpan-OptiPlex-7000:~/code/projects$ python pytorch/tools/dynamo/verify_dynamo.py   
pytorch/tools/dynamo/verify_dynamo.py:67: UserWarning: CUDA version mismatch, `torch` version: 11.7, env version: 11.8  
  warnings.warn(  
Python version: 3.8.15  
`torch` version: 1.14.0.dev20221202+cu117  
CUDA version: 11.8  
.. 忽略一些警告  
All required checks passed

官方提供的pytorch/tools/dynamo/verify_dynamo.py会检查2.0的新功能能否正常使用，如果All required checks passed说明就没问题，一共检查这几项，说实话，这几项我也没看明白，后续有机会慢慢看...

picture.image

然后跑几个模型试试，按照以下的方式进行测试，调用compiled_model = torch.compile(model)来对模型进行优化，对比优化前和优化后的速度以及余弦相似度：


      
        
import torch  
import time  
import torch._dynamo  
import torchvision.models as models  
  
torch._dynamo.config.verbose=True  
torch._dynamo.config.suppress_errors = True  
torch.set_float32_matmul_precision('high')  
  
model = models.resnet50().cuda()  
  
print("prepare model and input")  
dummy_input = torch.randn(1,3,1024,1024).cuda()  
  
NITER = 300  
print("warm...")  
for _ in range(10):  
    res = model(dummy_input)  
    torch.cuda.synchronize()  
      
print("begin eval ...")  
torch.cuda.synchronize()  
s = time.time()  
for _ in range(NITER):  
    res = model(dummy_input)  
    torch.cuda.synchronize()  
      
print('benchmark time (CUDA normal) (ms/iter)', (time.time() - s) / NITER * 1000)  
  
compiled_model = torch.compile(model)  
print("warm...")  
for _ in range(10):  
    res_compiled = compiled_model(dummy_input)  
    torch.cuda.synchronize()  
      
print("begin eval ...")  
torch.cuda.synchronize()  
s = time.time()  
for _ in range(NITER):  
    res_compiled = compiled_model(dummy_input)  
    torch.cuda.synchronize()  
      
print('benchmark time (torch.compiled) (ms/iter)', (time.time() - s) / NITER * 1000)  
print("check res cosine\_similarity")  
assert (  
    torch.nn.functional.cosine_similarity(  
        res.flatten(), res_compiled.flatten(), dim=0  
    )  
    > 0.9999  
)

测试结果如下，输入都是torch.randn(1,3,1024,1024).cuda()，其中reduce-overhead和max-autotune为torch.compile函数中的优化参数，除了res50，我也试了下dla结构的模型，其中with_dcn是带了自定义cuda实现的pytorch-dcn算子：

模型	无加速	compile-default	reduce-overhead	max-autotune
resnet50	26.6ms	22.2ms	22.1ms	报错
dla34_without_dcn	30.5ms	24.8ms	24.3ms	26.9ms
dla34_with_dcn	50.3ms	48.2ms	49.6ms	报错

compile的优化细节还没有细看，可以简单看下官方的描述：


      
        
# API NOT FINAL  
# default: optimizes for large models, low compile-time  
#          and no extra memory usage  
torch.compile(model)  
  
# reduce-overhead: optimizes to reduce the framework overhead  
#                and uses some extra memory. Helps speed up small models  
torch.compile(model, mode="reduce-overhead")  
  
# max-autotune: optimizes to produce the fastest model,  
#               but takes a very long time to compile  
torch.compile(model, mode="max-autotune")

reduce-overhead的意思是适合小模型，而max-autotune则是相当于trt或者tvm那样对整个模型进行编译优化了，选用这个的时候，compile的明显时间变长，不过效果嘛，我测试的几个模型没有明显提升甚至会编译失败，原因应该比较复杂，需要深入去看。

放一个torch.compile的函数介绍：


      
        
def compile(model: Optional[Callable] = None, *,  
            fullgraph: builtins.bool = False,  
            dynamic: builtins.bool = False,  
            backend: Union[str, Callable] = "inductor",  
            mode: Union[str, None] = None,  
            passes: Optional[Dict[str, Union[str, builtins.int, builtins.bool]]] = None,  
            **kwargs) -> Callable:  
    """  
    Optimizes given model/function using Dynamo and specified backend  
  
    Args:  
       model (Callable): Module/function to optimize  
       fullgraph (bool): Whether it is ok to break model into several subgraphs  
       dynamic (bool): Use dynamic shape tracing  
       backend (str or Callable): backend to be used  
       mode (str): Can be either "default", "reduce-overhead" or "max-autotune"  
       passes (dict): A dictionary of passes to the backend. Passes currently recognized by inductor backend:  
                       - static-memory  
                       - matmul-tune  
                       - matmul-padding  
                       - triton-autotune  
                       - triton-bmm  
                       - triton-mm  
                       - triton-convolution  
                       - rematerialize-threshold  
                       - rematerialize-acc-threshold  
  
    Example::  
  
        @torch.compile(passes={"matmul-padding": True}, fullgraph=True)  
        def foo(x):  
            return torch.sin(x) + torch.cos(x)  
  
    """

后续

细节比较多，我自己也有些疑问，不确定的点，也就先不写了hhh。

之后Pytorch2.0的版本，可能Dynamo会成为主要的模型的parser，替代torch.fx.trace和torch.jit.trace成为我们常用的工具。之后Pytorch转trt、转tvm、转乱七八槽一切也会慢慢往这个上头靠：

picture.image

总之，Pytorch这盘棋下的很大很大，对于我们来说，要学习的东西更多了，哎，又得卷了（fx-qat和triton后续文章在写了在写了，想写的东西太多，并行不过来了555）。

picture.image

参考资料

[1] 介绍: https://pytorch.org/get-started/pytorch-2.0/#technology-overview

往期回顾

关注 “ oldpan博客 ”，持续酝酿深度质量文

我是老潘，我们下期见~

   打上

星标 ✨不再错过老潘的及时推文

如果觉得有收获，来个

点赞加好看