開始使用¶
在閱讀本節之前，請務必閱讀 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 執行的最重要優化是融合 (fusion)。在上面的例子中，我們可以將 2 個讀取 (x, a) 和 2 個寫入 (a, b) 轉換為 1 個讀取 (x) 和 1 個寫入 (b)，這對於記憶體頻寬（您可以將數據發送到 GPU 的速度）是瓶頸的新型 GPU 尤其重要，而非計算（您的 GPU 可以處理浮點運算的速度）。
Inductor 提供的另一個主要優化是自動支援 CUDA graphs。CUDA graphs 有助於消除從 Python 程式啟動個別 kernels 的開銷，這對於較新的 GPU 尤其重要。
TorchDynamo 支援許多不同的後端，但 TorchInductor 專門透過產生 Triton kernels 來工作。讓我們將上面的例子儲存到名為 example.py 的檔案中。我們可以透過執行 TORCH_COMPILE_DEBUG=1 python example.py 來檢查產生的 Triton kernels 的程式碼。當腳本執行時，您應該會在終端機上看到 DEBUG 訊息。在日誌接近尾聲時，您應該會看到一個資料夾的路徑，其中包含 torchinductor_<您的使用者名稱>。在該資料夾中，您可以找到 output_code.py 檔案，其中包含類似於以下的產生的 kernel 程式碼
@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 kernel 中，並且暫時變數保存在具有非常快速存取的暫存器中。
在此處閱讀更多關於 Triton 效能的資訊：這裡。由於程式碼是用 Python 編寫的，因此即使您沒有編寫過那麼多的 CUDA kernels，也很容易理解。
接下來，讓我們嘗試一個真實的模型，例如來自 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)
如果您從模型和 encoded_input 中刪除 to(device="cuda:0")，那麼 Triton 將產生 C++ kernels，這些 kernels 將針對在您的 CPU 上執行進行優化。您可以檢查 Triton 或 C++ kernels 以用於 BERT。它們比我們上面嘗試的三角函數例子更複雜，但您可以類似地瀏覽它，看看您是否了解 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 的工作原理有了基本的了解。以下是您接下來可以查看的內容
