入門#

創建於：2025 年 6 月 16 日 | 最後更新於：2025 年 6 月 16 日

在閱讀本節之前，請務必閱讀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()。在 eager 模式下，逐元素運算子的效率不高，因為每個運算子都需要從記憶體中讀取一個張量，進行一些更改，然後將這些更改寫回。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_<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.tw/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)

如果您從模型和 encoded_input 中刪除 to(device="cuda:0")，那麼 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 的工作原理有了基本的瞭解。接下來您可以檢視以下內容：

入門#

使用預訓練模型#

後續步驟#

文件

教程

資源