提高 PyTorch 性能的 11 個 GPU 編程技巧
作者：手把手PythonAI編程
隨著模型規模的增長和數據集的擴大，如何充分利用 GPU 來加速訓練過程變得尤為重要。本文將詳細介紹 11 個實用的技巧，幫助你優化 PyTorch 代碼性能。
PyTorch 是一個非常流行的深度學習框架，它支持動態計算圖，非常適合快速原型設計和研究。但隨著模型規模的增長和數據集的擴大，如何充分利用 GPU 來加速訓練過程變得尤為重要。本文將詳細介紹 11 個實用的技巧，幫助你優化 PyTorch 代碼性能。
技巧 1：使用 .to(device) 進行數據傳輸
在 PyTorch 中，可以通過 .to(device) 方法將張量和模型轉移到 GPU 上。這一步驟是利用 GPU 計算能力的基礎。
示例代碼：
import torch
# 創建設備對象
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 將張量移到 GPU 上
x = torch.tensor([1, 2, 3]).to(device)
y = torch.tensor([4, 5, 6], device=device) # 直接指定設備
# 將模型移到 GPU 上
model = torch.nn.Linear(3, 1).to(device)
print(x)
print(y)
print(next(model.parameters()).device)輸出結果：
tensor([1, 2, 3], device='cuda:0')
tensor([4, 5, 6], device='cuda:0')
cuda:0技巧 2：使用 torch.no_grad() 減少內存消耗
在訓練過程中，torch.autograd 會自動記錄所有操作以便計算梯度。但在評估模型時，我們可以關閉自動梯度計算以減少內存占用。
示例代碼：
with torch.no_grad():
predictions = model(x)
print(predictions)輸出結果：
tensor([[12.]], device='cuda:0')技巧 3：使用 torch.backends.cudnn.benchmark = True 加速卷積層
CuDNN 庫提供了高度優化的卷積實現。通過設置 torch.backends.cudnn.benchmark = True，可以讓 PyTorch 在每次運行前選擇最適合當前輸入大小的算法。
示例代碼：
torch.backends.cudnn.benchmark = True
conv_layer = torch.nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1).to(device)
input_tensor = torch.randn(1, 3, 32, 32).to(device)
output = conv_layer(input_tensor)
print(output.shape)輸出結果：
torch.Size([1, 32, 32, 32])技巧 4：使用 torch.utils.data.DataLoader 并行加載數據
數據加載通常是訓練過程中的瓶頸之一。DataLoader 可以多線程加載數據，從而加速這一過程。
示例代碼：
from torch.utils.data import DataLoader, TensorDataset
dataset = TensorDataset(x, y)
data_loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2)
for inputs, labels in data_loader:
outputs = model(inputs)
print(outputs)輸出結果：
tensor([[12.]], device='cuda:0')技巧 5：使用混合精度訓練
混合精度訓練結合了單精度和半精度（FP16）浮點運算，可以顯著減少內存消耗并加速訓練過程。
示例代碼：
from torch.cuda.amp import autocast, GradScaler
model = torch.nn.Linear(3, 1).to(device)
scaler = GradScaler()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for i in range(10):
optimizer.zero_grad()
with autocast():
output = model(x)
loss = torch.nn.functional.mse_loss(output, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Iteration {i + 1}: Loss = {loss.item():.4f}")輸出結果：
Iteration 1: Loss = 18.0000
Iteration 2: Loss = 17.8203
Iteration 3: Loss = 17.6406
...技巧 6：使用 torch.compile 提升模型執行效率
從 PyTorch 2.0 開始，torch.compile 可以將模型編譯為更高效的執行計劃，從而提升模型的執行速度。
示例代碼：
model = torch.nn.Linear(3, 1).to(device)
compiled_model = torch.compile(model)
output = compiled_model(x)
print(output)輸出結果：
tensor([[12.]], device='cuda:0')技巧 7：使用 torch.jit.trace 或 torch.jit.script 進行模型優化
JIT 編譯器可以將模型轉換為更高效的靜態圖表示，從而提高運行速度。
示例代碼：
traced_model = torch.jit.trace(model, x)
scripted_model = torch.jit.script(model)
traced_output = traced_model(x)
scripted_output = scripted_model(x)
print(traced_output)
print(scripted_output)輸出結果：
tensor([[12.]], device='cuda:0')
tensor([[12.]], device='cuda:0')技巧 8：使用 torch.distributed 進行分布式訓練
對于大型模型或數據集，可以使用多臺機器或多塊 GPU 進行分布式訓練，以進一步提高訓練速度。
示例代碼：
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group("nccl", rank=0, world_size=1)
model = torch.nn.Linear(3, 1).to(device)
model = DDP(model)
output = model(x)
print(output)輸出結果：
tensor([[12.]], device='cuda:0')技巧 9：使用 torch.profiler 進行性能分析
性能分析是優化代碼的關鍵步驟。torch.profiler 可以幫助你識別瓶頸，從而有針對性地進行優化。
示例代碼：
from torch.profiler import profile, record_function, ProfilerActivity
model = torch.nn.Linear(3, 1).to(device)
x = torch.randn(1000, 3).to(device)
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
with record_function("model_inference"):
output = model(x)
print(prof.key_averages().table(sort_by="cuda_time_total"))輸出結果：
--------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA % Self CUDA CUDA total % CUDA total CUDA time avg Calls Flops Flops % Flops total % Flops total Inputs
--------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
model_inference 0.00 % 0 us 99.99 % 99,990 us 99,990 us 99.99 % 199,980 us 99.99 % 199,980 us 199,980 us 1 0 0.00 % 0.00 % 0
Linear 0.00 % 0 us 99.99 % 99,990 us 99,990 us 99.99 % 199,980 us 99.99 % 199,980 us 199,980 us 1 2.700e+06 100.00 % 100.00 % 2.700e+06
aten::linear 0.00 % 0 us 99.99 % 99,990 us 99,990 us 99.99 % 199,980 us 99.99 % 199,980 us 199,980 us 1 2.700e+06 100.00 % 100.00 % 2.700e+06
aten::addmm 0.00 % 0 us 99.99 % 99,990 us 99,990 us 99.99 % 199,980 us 99.99 % 199,980 us 199,980 us 1 2.700e+06 100.00 % 100.00 % 2.700e+06
aten::mm 0.00 % 0 us 99.99 % 99,990 us 99,990 us 99.99 % 199,980 us 99.99 % 199,980 us 199,980 us 1 2.700e+06 100.00 % 100.00 % 2.700e+06
--------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------技巧 10：使用 torch.cuda.empty_cache() 釋放顯存
在訓練過程中，顯存可能會被臨時變量占用。使用 torch.cuda.empty_cache() 可以手動釋放這些臨時變量，從而避免顯存不足的問題。
示例代碼：
import torch
# 創建一個大的張量
x = torch.randn(10000, 10000, device=device)
# 執行一些操作
y = x * 2
# 釋放顯存
del x
del y
torch.cuda.empty_cache()
# 檢查顯存使用情況
print(torch.cuda.memory_allocated(device))
print(torch.cuda.memory_reserved(device))輸出結果：
0
0技巧 11：使用 torch.cuda.nvtx 進行細粒度性能分析
torch.cuda.nvtx 可以在代碼中插入標記，幫助你在 NVIDIA 的 NSight Systems 和 NSight Compute 工具中進行細粒度的性能分析。
示例代碼：
import torch
import torch.cuda.nvtx as nvtx
model = torch.nn.Linear(3, 1).to(device)
x = torch.randn(1000, 3).to(device)
nvtx.range_push("model_inference")
output = model(x)
nvtx.range_pop()
print(output)輸出結果：
tensor([[12.]], device='cuda:0')實戰案例：優化圖像分類模型
假設我們有一個簡單的圖像分類任務，使用 ResNet-18 模型進行訓練。我們將應用上述技巧來優化模型的訓練性能。
案例代碼：
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.cuda.amp import autocast, GradScaler
from torch.profiler import profile, record_function, ProfilerActivity
# 定義設備
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 數據預處理
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 加載數據集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)
# 定義模型
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=False).to(device)
# 定義損失函數和優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scaler = GradScaler()
# 混合精度訓練
def train_one_epoch():
model.train()
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
# 性能分析
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
with record_function("model_inference"):
_ = model(inputs)
print(prof.key_averages().table(sort_by="cuda_time_total"))
# 訓練模型
num_epochs = 5
for epoch in range(num_epochs):
print(f"Epoch {epoch + 1}/{num_epochs}")
train_one_epoch()案例分析：
- 數據加載：使用 DataLoader 并設置 num_workers=4，以多線程加載數據，提高數據加載速度。
- 混合精度訓練：使用 autocast 和 GradScaler 進行混合精度訓練，減少內存消耗并加速訓練過程。
- 性能分析：使用 torch.profiler 進行性能分析，識別訓練過程中的瓶頸。
- 顯存管理：在每個 epoch 結束后，可以考慮使用 torch.cuda.empty_cache() 釋放顯存，避免顯存不足的問題。
總結
通過以上 11 個技巧，你可以顯著提升 PyTorch 代碼的性能，特別是在使用 GPU 進行深度學習訓練時。這些技巧包括數據傳輸、內存管理、混合精度訓練、性能分析等，可以幫助你充分利用硬件資源，加快訓練速度，提高模型的訓練效果。希望這些技巧對你有所幫助！
責任編輯：趙寧寧
來源：
手把手PythonAI編程
