提高 PyTorch 性能的 11 个 GPU 编程技巧-51CTO.COM

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 进行深度学习训练时。这些技巧包括数据传输、内存管理、混合精度训练、性能分析等，可以帮助你充分利用硬件资源，加快训练速度，提高模型的训练效果。希望这些技巧对你有所帮助！