编译器优化那些事儿（21）：torch.compile 命令加速模型

文章预览

在 PyTorch 2.0 的重大更新中， torch.compile() 或许是其中最引人注目的一项。对于各种 PyTorch 模型， 使用 torch.compile 调用即可实现加速，在各种数据集上都获得了显著的加速效果（如图1所示）。这样的效果是如何实现的呢？接下来我们将一探究竟。 图1 torch.compile在A100 GPU上相比eager模式的加速效果 1. 背景：PyTorch 的强与弱 在众多现代机器学习框架中，PyTorch 凭借 eager 求值 和 Python 原生 两大特性而广受研究者欢迎。eager 求值使得 PyTorch 应用易于调试，而 Python 原生为扩展模型提供了巨大便利。 然而，在易调试和易扩展的另一面，这些特性也带来了性能上的副作用。相比先定义再运行的静态机器学习框架（如：Tensorflow），长期以来 PyTorch 饱受性能不强的困扰。eager 求值导致框架每次只能看到单个算子执行，缺乏对于整个模型的全局视角，难以通过 ………………………………