PyTorch 2.0正式发布

12月2日，PyTorch 2.0正式发布！

这次的更新不仅将PyTorch的性能推到了新的高度，同时也加入了对动态形状和分布式的支持。

此外，2.0系列还会将PyTorch的部分代码从C++移回Python。

目前，PyTorch 2.0还处在测试阶段，预计第一个稳定版本会在2023年3月初面世。

PyTorch 2.x：更快、更Python！

在过去的几年里，PyTorch从1.0到最近的1.13进行了创新和迭代，并转移到新成立的PyTorch基金会，成为Linux基金会的一部分。

当前版本的PyTorch所面临的挑战是，eager-mode难以跟上不断增长的GPU带宽和更疯狂的模型架构。

而PyTorch 2.0的诞生，将从根本上改变和提升了PyTorch在编译器级别下的运行方式。众所周知，PyTorch中的（Py）来自于数据科学中广泛使用的开源Python编程语言。

然而，PyTorch的代码却并没有完全采用Python，而是把一部分交给了C++。

不过，在今后的2.x系列中，PyTorch项目团队计划将与torch.nn有关的代码移回到Python中。

除此之外，由于PyTorch 2.0是一个完全附加的（和可选的）功能，因此2.0是100%向后兼容的。

也就是说，代码库是一样的，API也是一样的，编写模型的方式也是一样的。

更多的技术支持

使用Python框架评估钩子安全地捕获PyTorch程序，这是团队5年来在graph capture方面研发的一项重大创新。

重载了PyTorch的autograd引擎，作为一个追踪的autodiff，用于生成超前的反向追踪。

将约2000多个PyTorch运算符归纳为约250个原始运算符的封闭集，开发人员可以针对这些运算符构建一个完整的PyTorch后端。大大降低了编写PyTorch功能或后端的障碍。

一个深度学习编译器，可以为多个加速器和后端生成快速代码。对于英伟达的GPU，它使用OpenAI Triton作为关键构建模块。

值得注意的是，TorchDynamo、AOTAutograd、PrimTorch和TorchInductor都是用Python编写的，并支持动态形状。

更快的训练速度

通过引入新的编译模式「torch.compile」，PyTorch 2.0用一行代码，就可以加速模型的训练。

这里不用任何技巧，只需运行torch.compile()即可，仅此而已：

opt_module = torch.compile(module)

为了验证这些技术，团队精心打造了测试基准，包括图像分类、物体检测、图像生成等任务，以及各种NLP任务，如语言建模、问答、序列分类、推荐系统和强化学习。其中，这些基准可以分为三类：

测试结果表明，在这163个跨越视觉、NLP和其他领域的开源模型上，训练速度得到了38%-76%的提高。

在NVIDIA A100 GPU上的对比

此外，团队还在一些流行的开源PyTorch模型上进行了基准测试，并获得了从30%到2倍的大幅加速。

开发者Sylvain Gugger表示：「只需添加一行代码，PyTorch 2.0就能在训练Transformers模型时实现1.5倍到2.0倍的速度提升。这是自混合精度训练问世以来最令人兴奋的事情！」

本文转载于公众号新智元，更多内容可查看

内容中包含的图片若涉及版权问题，请及时与我们联系删除