入门指南
在阅读本节之前，请确保已经阅读了 torch.compiler。
让我们从查看一个简单的 torch.compile 示例开始，该示例演示了如何使用 torch.compile 进行推理。此示例演示了 torch.cos() 和 torch.sin() 功能，它们是点操作符的示例，因为它们逐元素对向量进行操作。此示例可能不会显示显著的性能提升，但应该有助于您形成对如何在您的程序中使用 torch.compile 的直观理解。
注意
要运行此脚本，您需要在您的机器上至少有一个 GPU。如果您没有 GPU，可以删除以下片段中的 .to(device="cuda:0") 代码，它将在 CPU 上运行。您还可以设置设备为 xpu:0 以在 Intel® GPU 上运行。
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() 的东西。在急切模式下，逐点运算是不理想的，因为每个运算都需要从内存中读取一个张量，进行一些更改，然后再将更改写回。感应器执行的最重要优化是融合。在上面的例子中，我们可以将 2 次读取（ x ， a ）和 2 次写入（ a ， b ）转换为 1 次读取（ x ）和 1 次写入（ b ），这对于内存带宽（将数据发送到 GPU 的速度）是瓶颈的新 GPU 尤其重要，而不是计算（GPU 执行浮点运算的速度）。
感应器提供的另一个主要优化是自动支持 CUDA 图。CUDA 图有助于消除从 Python 程序启动单个内核的开销，这对于新 GPU 尤其相关。
TorchDynamo 支持许多不同的后端，但 TorchInductor 特别通过生成 Triton 内核来实现。让我们将上面的示例保存到名为 example.py 的文件中。我们可以通过运行 TORCH_COMPILE_DEBUG=1 python example.py 来检查生成的 Triton 内核代码。当脚本执行时，你应该会在终端看到 DEBUG 消息。在日志的末尾附近，你应该会看到一个包含 torchinductor_<your_username> 的文件夹的路径。在那个文件夹中，你可以找到包含类似以下内容的生成内核代码的 output_code.py 文件：
@pointwise(size_hints=[16384], filename=__file__, triton_meta={'signature': {'in_ptr0': '*fp32', 'out_ptr0': '*fp32', 'xnumel': 'i32'}, 'device': 0, 'constants': {}, 'mutated_arg_names': [], 'configs': [AttrsDescriptor(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, other=0.0)
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
# Copy pasted from here 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") # This is the only line of code that we changed
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
output = model(**encoded_input)
如果您从模型中移除 to(device="cuda:0") 和 encoded_input ，那么 Triton 将生成针对您的 CPU 运行的 C++内核，您可以对 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 的工作原理有了基本的了解。接下来您可以检查以下内容：
