入门指南¶

在阅读本节之前，请确保阅读了torch.compiler。

让我们从一个简单的 torch.compile 示例开始，该示例展示了如何使用 torch.compile 进行推理。这个示例展示了 torch.cos() 和 torch.sin() 功能，它们是逐点操作符的示例，因为它们逐元素地对向量进行操作。这个示例可能不会显示出显著的性能提升，但应该有助于你形成对如何在自定义程序中使用 torch.compile 的直观理解。

注意

要运行此脚本，您的机器上至少需要有一块GPU。如果您没有GPU，可以删除下面的代码片段中的.to(device="cuda:0")代码，它将在CPU上运行。

import torch
def fn(x):
   a = torch.cos(x)
   b = torch.sin(a)
   return b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor)

一个你可能更想使用的更著名的逐点操作符是类似于 torch.relu() 的东西。在急切模式下的逐点操作是次优的，因为每个操作都需要从内存中读取一个张量，进行一些更改，然后将这些更改写回。Inductor 执行的最重要的优化是融合。在上面的例子中，我们可以将 2 次读取（x，a）和 2 次写入（a，b）转换为 1 次读取（x）和 1 次写入（b），这对于新型的 GPU 尤其重要，因为瓶颈在于内存带宽（你能多快地将数据发送到 GPU）而不是计算（你的 GPU 能多快地处理浮点运算）。

Inductor 提供的另一个主要优化是自动支持 CUDA 图。 CUDA 图有助于消除从 Python 程序启动单个内核的开销，这对于较新的 GPU 尤其重要。

TorchDynamo 支持许多不同的后端，但 TorchInductor 特别通过生成 Triton 内核来工作。让我们将上面的示例保存到一个名为 example.py 的文件中。我们可以通过运行 TORCH_COMPILE_DEBUG=1 python example.py 来检查生成的 Triton 内核代码。当脚本执行时，您应该会看到 DEBUG 消息打印到终端。在日志的末尾附近，您应该会看到一个包含 torchinductor_ 的文件夹路径。在该文件夹中，您可以找到包含生成的内核代码的 output_code.py 文件，类似于以下内容：

@pointwise(size_hints=[16384], filename=__file__, triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'mutated_arg_names': [], 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
   xnumel = 10000
   xoffset = tl.program_id(0) * XBLOCK
   xindex = xoffset + tl.arange(0, XBLOCK)[:]
   xmask = xindex < xnumel
   x0 = xindex
   tmp0 = tl.load(in_ptr0 + (x0), xmask)
   tmp1 = tl.cos(tmp0)
   tmp2 = tl.sin(tmp1)
   tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)

注意

上述代码片段是一个示例。根据您的硬件，您可能会看到生成的不同代码。

你可以验证融合 cos 和 sin 确实发生了，因为 cos 和 sin 操作发生在一个 Triton 内核中，并且临时变量保存在寄存器中，具有非常快的访问速度。

阅读更多关于Triton性能的信息这里。因为代码是用Python编写的，即使你没有编写过很多CUDA内核，也很容易理解。

接下来，让我们尝试一个真实的模型，如来自 PyTorch hub 的 resnet50。

import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(1,3,64,64))

而且这并不是唯一可用的后端，你可以在 REPL 中运行 torch.compiler.list_backends() 来查看所有可用的后端。接下来可以尝试 cudagraphs 作为灵感。

使用预训练模型¶

PyTorch用户经常利用来自 transformers 或 TIMM 的预训练模型，其中一个设计目标是TorchDynamo和TorchInductor能够开箱即用地支持人们想要创建的任何模型。

让我们直接从HuggingFace hub下载一个预训练模型并进行优化：

import torch
from transformers import BertTokenizer, BertModel
# 从这里复制粘贴 https://huggingface.co/bert-base-uncased
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
model = torch.compile(model, backend="inductor") # 这是我们唯一更改的代码行
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
output = model(**encoded_input)

如果你从模型和encoded_input中移除to(device="cuda:0")，那么Triton将生成C++内核，这些内核将针对在CPU上运行进行优化。你可以检查BERT的Triton或C++内核。它们比我们上面尝试的三角函数示例更复杂，但你可以类似地浏览它们，看看你是否理解PyTorch的工作原理。

同样地，让我们尝试一个TIMM示例：

import timm
import torch
model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(64,3,7,7))

下一步¶

在本节中，我们回顾了一些推理示例，并初步了解了torch.compile的工作原理。接下来你可以查看以下内容：