更快的Pytorch，2.X新时代到来

更快的Pytorch 2.X已经到来

更快的Pytorch 2.X

Pytorch官方团队宣布Pyotch 2.X的新时代已经到来，新版本将会引入torch.compile，这是一种编译模式，可以在不更改模型代码的情况下加速模型。并且torch.compile是一个可选特性，因此Pyotch 2.0 是 100% 向后兼容的。

研发团队在各种机器学习领域测试了 163 个开源模型，torch.compile在 93% 的情况下都有效，模型在 NVIDIA A100 GPU 上的训练速度提高了 43% 。 在 Float32 精度下，平均运行速度提高了 21% ，而在AMP精度下，平均运行速度提高了 51% 。

torch.compile 现还处于开发的早期阶段。预计将在 2023 年 3 月上旬发布第一个稳定的 2.0 版本。

使用方式

Pyotch 2.X引入了一个简单的函数 torch.compile 来包装模型, 并返回一个编译后的模型。

compiled_model = torch.compile(model)

返回的 compiled_model 包含对模型的引用，并将前向函数编译为更优化的版本。在编译模型时，会给出了几用来调整它的设置，如下所示：

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

• mode: 指定编译器在编译时应该优化什么。默认模式是一种预设，它会尝试在不花费太长时间编译或使用额外内存的情况下高效编译。

• dynamic: 指定是否启用动态形状的代码路径。某些编译器优化不能应用于动态整形程序。明确表示您想要一个具有动态形状还是静态形状的编译程序将有助于编译器为您提供更好的优化代码。

• fullgraph: 类似于 Numba 的 nopython。它将整个程序编译成一个图形，或者给出一个错误来解释为什么它不能这样做。大多数用户不需要使用此模式。如果您非常注重性能，那么您可以尝试使用它。

• backend: 指定要使用的编译器后端。默认情况下，使用 TorchInductor。

编译旨在在默认模式下提供最大的好处和最大的灵活性。下面，让我们使用一个随机数据，编译真实模型并运行它的完整示例：

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()

第一次运行compiled_model(x)时，它会编译模型。 因此，运行需要更长的时间。 随后的运行速度会很快。

torch.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")

安装

对于GPU（新一代GPU的性能会大大提高）

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

对于CPU

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

欢迎加入我们

欢迎加入我们的QQ交流群获取使用的数据：837977579

欢迎关注我们的微信公众号获取更多内容

今天的分享就到这里了，敬请期待下一篇!

最后欢迎大家分享转发，您的点赞是对我的鼓励和肯定！

往期回顾

参考资料

1:入门教程

2:文章原文