入门¶
让我们从一个简单的例子开始。请注意，你的 GPU 越新，可能会看到更显著的速度提升。
import torch
def fn(x, y):
a = torch.cos(x).cuda()
b = torch.sin(y).cuda()
return a + b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor, input_tensor)
这个示例实际上不会运行得更快。它的目的是演示
torch.cos() 和 torch.sin() 功能，它们是逐点操作的示例，因为它们对向量中的每个元素进行操作。你可能想要使用的一个更著名的逐点操作是类似 torch.relu() 的操作。在即时模式下，逐点操作效率不高，因为每个操作都需要从内存中读取张量，进行一些修改，然后再将这些修改写回内存。Inductor 所做的最重要的优化是融合。因此回到我们的示例，我们可以将两次读取和两次写入合并为一次读取和一次写入，这对于新 GPU 来说尤其关键，因为瓶颈是内存带宽（数据发送到 GPU 的速度）而不是计算能力（GPU 执行浮点运算的速度）。
Inductor 提供的另一项重要优化是自动支持 CUDA 图（CUDA graphs）。 CUDA 图有助于消除从 Python 程序中启动单个内核所产生的开销，这对于较新的 GPU 来说尤为重要。
TorchDynamo 支持许多不同的后端，但 inductor 特别是通过生成 Triton 内核来工作的
并且我们可以通过运行 TORCH_COMPILE_DEBUG=1 python trig.py 来检查它们
实际生成的内核是
@pointwise(size_hints=[16384], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
@triton.jit
def kernel(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
xnumel = 10000
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])
xmask = xindex < xnumel
x0 = xindex
tmp0 = tl.load(in_ptr0 + (x0), xmask)
tmp1 = tl.sin(tmp0)
tmp2 = tl.sin(tmp1)
tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)
你可以验证将两个 sin 合并确实发生了，
因为这两个 sin 操作发生在单个 Triton 内核中，
临时变量存储在寄存器中，访问速度非常快。
你可以在这里阅读更多关于Triton性能的信息 这里 但关键是它在Python中 所以即使你没有编写很多CUDA内核，也很容易理解。
接下来，让我们尝试一个真正的模型，比如从 PyTorch hub 中的 resnet50。
import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
opt_model = torch.compile(model, backend="inductor")
model(torch.randn(1,3,64,64))
And that is not the only available backend, you can run in a REPL
torch._dynamo.list_backends() to see all the available backends. Try out the
cudagraphs or nvfuser next as inspiration.
现在让我们做一些更有趣的事情，我们的社区经常使用 transformers 或 TIMM 中的预训练模型，而我们的设计目标之一是让 Dynamo 和 inductor 能够开箱即用，支持用户想要编写的任何模型。
因此，我们将直接从 HuggingFace 仓库下载一个预训练模型并对其进行优化：
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 将会生成 C++ 内核，这些内核将被优化为在你的 CPU 上运行。你可以检查 BERT 的 Triton 或 C++ 内核，它们显然比我们上面的三角函数示例要复杂得多，但如果你理解 PyTorch 的话，也可以同样快速浏览并理解它们。
同样，我们来尝试一个 TIMM 示例
import timm
import torch._dynamo as dynamo
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))
我们使用Dynamo和inductor的目标是构建一个覆盖范围最广的机器学习编译器，它应该能够处理你投给它的任何模型。
现有后端¶
TorchDynamo 拥有一个不断增长的后端列表，可以在
backends 文件夹中找到
或 torch._dynamo.list_backends() 每个后端及其可选依赖项。
一些常用的后端包括：
- Training & inference backends:
- Inference-only backends:
为什么你需要另一种优化PyTorch代码的方法？¶
虽然 PyTorch 生态系统中存在其他一些代码优化工具，但每个工具都有其自己的流程。 以下是一些现有方法及其局限性的例子：
torch.jit.trace()在无法追踪时会静默出错，例如： 在控制流期间torch.jit.script()需要对用户代码或库代码进行修改，通过添加类型注解并删除非 PyTorch 代码torch.fx.symbolic_trace()要么正确追踪，要么给出一个硬错误，但它仅限于可追踪的代码，因此仍然无法处理控制流torch._dynamo可以直接工作并生成部分图。 它仍然有选项可以生成一个单一的图，其中包含nopython=True适用于某些情况 但允许更平滑的过渡，在不需要代码修改的情况下优化部分图