如何提高-PyTorch-训练循环的效率

如何提高 PyTorch 训练循环的效率

原文：towardsdatascience.com/improve-efficiency-of-your-pytorch-training-loop/

训练深度学习模型不仅仅是将数据提交给反向传播算法。通常，决定项目成功或失败的关键因素在于一个不太引人注目但绝对关键的领域：数据管道的效率。

效率低下的训练基础设施浪费时间、资源和金钱，导致图形处理单元（GPU）空闲，这种现象称为 GPU 饥饿。这种低效不仅延迟了开发，还增加了运营成本，无论是在云基础设施还是在本地基础设施上。

本文旨在作为识别和解决 PyTorch 训练周期中最常见瓶颈的实用和基本指南。

分析将集中在数据管理上，这是每个训练循环的核心，并将展示如何通过有针对性的优化释放硬件的潜力，从理论方面到实际实验。

总结来说，通过阅读这篇文章，你将学到：

• 减缓神经网络开发和训练的常见瓶颈

• 在 PyTorch 中优化训练循环的基本原则

• 训练中的并行性和内存管理

---

训练优化的动机

提高深度学习模型的训练效率是一项战略性的必要措施——它直接转化为成本和计算时间的显著节省。

更快的训练允许：

• 更快的测试周期

• 新想法的验证

• 探索不同的架构和微调超参数

这加速了模型的生命周期，使组织能够更快地进行创新并将解决方案推向市场。

例如，训练优化允许公司快速分析大量数据以识别趋势和模式，这对于制造中的模式识别或预测性维护是一项关键任务。

最常见瓶颈的分析

减缓通常表现为 CPU、GPU、内存和存储设备之间的复杂交互。

这里列出了可能导致神经网络训练变慢的主要瓶颈：

• I/O 和数据： 主要问题是 GPU 饥饿，即 GPU 空闲等待 CPU 加载和预处理下一批数据。这在无法完全加载到 RAM 的大型数据集中很常见。磁盘速度至关重要：NVMe SSD 的速度可以比传统 HDD 快 35 倍。

• GPU: 当 GPU 过载（计算密集型模型）或更常见的是，由于 CPU 提供的数据不足而未充分利用时发生。与擅长顺序处理的 CPU 不同，GPU 具有众多低速核心，优化了并行处理。

• 内存：内存耗尽，通常表现为著名的RuntimeError: CUDA out of memory错误，迫使减少批大小。梯度堆叠技术可以模拟更大的批大小，但它不会增加吞吐量。

为什么 CPU 和 I/O 经常是主要限制？

优化的一个关键方面是理解“级联瓶颈”。

在典型的训练系统中，GPU 是计算引擎，而 CPU 负责数据准备。如果磁盘速度慢，CPU 大部分时间都在等待数据，成为主要的瓶颈。因此，没有数据处理的 GPU 保持空闲。

这种行为导致了一个错误的信念，即问题出在 GPU 硬件上，而实际上不效率的是数据供应链。在不解决上游瓶颈的情况下增加 GPU 处理能力是浪费时间，因为训练性能永远不会超过系统中最慢的组件。因此，有效优化的第一步是识别和解决根本问题，这通常在于 I/O 或数据管道。

分析和优化的工具和库

有效的优化需要数据驱动的方法，而不是试错。PyTorch 提供了诊断瓶颈并改进训练周期的工具和原语。以下是我们的实验的三个关键要素：

• Dataset 和 Dataloader

• TorchVision

• 分析器

PyTorch 中的 Dataset 和 Dataloader

高效的数据管理是任何训练循环的核心。PyTorch 提供了两个基本抽象，称为Dataset 和 Dataloader。

这里有一个简要概述

• torch.utils.data.Dataset

  这是表示一组样本及其标签的基础类。

  要创建自定义数据集，只需实现三个方法：

  • __init__: 初始化路径或数据连接，

  • __len__: 返回数据集的长度，

  • __getitem__: 加载并可选地转换单个样本。

• torch.utils.data.DataLoader

  它是包装数据集并使其高效可迭代的接口。

  它自动处理：

  • 批处理（batch_size），

  • 重新洗牌（shuffle=True），

  • 并行加载（num_workers），

  • 内存管理（pin_memory）

TorchVision：计算机视觉的标准数据集和操作

TorchVision 是 PyTorch 的计算机视觉领域库，旨在加速原型设计和基准测试。

它的主要实用工具包括：

• 预定义数据集: CIFAR-10、MNIST、ImageNet 等，已作为Dataset的子类实现。非常适合快速测试，无需构建自定义数据集。

• 常见转换: 缩放、归一化、旋转、数据增强。这些操作可以通过transforms.Compose组合，并在加载时即时执行，从而减少手动预处理。

• 预训练模型：适用于分类、检测和分割任务，可作为基线或用于迁移学习。

示例：

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5])
])

train_data = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)

PyTorch Profiler：性能诊断工具

PyTorch Profiler 允许您精确了解您的执行时间在 CPU 和 GPU 上的具体花费。

关键特性：

• CUDA 操作符和内核的详细分析。

• 多设备支持（CPU/GPU）。

• 以.json交互格式导出结果或使用 TensorBoard 进行可视化。

示例：

import torch
import torch.profiler as profiler

def train_step(model, dataloader, optimizer, criterion):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

with profiler.profile(
    activities=[profiler.ProfilerActivity.CPU, 
    profiler.ProfilerActivity.CUDA],
    on_trace_ready=profiler.tensorboard_trace_handler("./log")
) as prof:

    train_step(model, dataloader, optimizer, criterion)

print(prof.key_averages().table(sort_by="cuda_time_total"))

训练周期的构建和分析

PyTorch 中的训练循环是一个迭代过程，对于每一批数据，重复一系列基本步骤来教会网络三个基本阶段：

1. 正向传播：模型从输入批次计算预测。PyTorch 在此阶段动态构建计算图（autograd），以跟踪操作并为梯度计算做准备。

2. 反向传播：反向传播计算损失函数相对于所有模型参数的梯度，使用链式法则。此过程通过调用loss.backward()触发。在每次反向传播之前，我们必须使用optimizer.zero_grad()重置梯度，因为 PyTorch 默认会累积它们。

3. 更新权重：优化器（torch.optim）使用计算出的梯度来更新模型权重，最小化损失。调用optimizer.step()执行当前批次的最终更新。

周期中可能出现各种减速点。如果从DataLoader加载批次较慢，GPU 将保持空闲。如果模型计算量很大，GPU 将饱和。CPU 和 GPU 之间的数据传输是另一个潜在的效率低下来源，表现为cudaMemcpyAsync分析操作的长执行时间。

训练瓶颈几乎永远不会是 GPU，而是导致其停机时间的数据管道中的低效。

主要目标是确保 GPU 永远不会饿死，保持数据供应的恒定。

优化利用了 CPU（适用于 I/O 和顺序处理）和 GPU（适用于并行计算）架构之间的对比。如果数据集太大而无法放入 RAM，基于 Python 的生成器可能成为训练复杂模型的一个重大障碍。

一个例子可能是一个训练循环，其中当 GPU 运行时，CPU 处于空闲状态，而当 CPU 运行时，GPU 处于空闲状态，如下所示：

该图像描绘了一个经典的不高效数据管理案例。图片由作者提供。

CPU 和 GPU 之间的批处理管理

优化过程基于重叠的概念：DataLoader使用多个工作者（num_workers > 0），在 CPU 上并行准备下一个批次，同时 GPU 处理当前的批次。

优化DataLoader确保 CPU 和 GPU 异步并发工作。如果一批次的预处理时间大约等于 GPU 的计算时间，理论上可以将训练过程的速度提高一倍。

这种预加载行为可以通过 DataLoader 的prefetch_factor参数进行控制，该参数确定每个工作者预加载的批次数量。

诊断瓶颈的方法

使用 PyTorch Profiler 有助于将优化过程转化为数据驱动的诊断。通过分析已过时间指标，您可以识别效率低下的根本原因：

Profiler 检测到的症状	诊断（瓶颈）	推荐的解决方案
DataLoader的Self CPU total %高	CPU 侧的预处理和数据加载缓慢	增加num_workers
cudaMemcpyAsync执行时间高	CPU 和 GPU 内存之间的数据传输缓慢	启用pin_memory=True

数据加载优化技术

在DataLoader中，PyTorch 实现的最有效技术是工作者并行和锁定内存（pinned_memory）的使用。

使用工作者实现并行

DataLoader中的num_workers参数启用多进程，创建子进程并行加载和预处理数据。这显著提高了数据加载吞吐量，有效地重叠了训练和下一批次的准备。

• 好处： 减少 GPU 等待时间，尤其是在处理大型数据集或复杂的预处理（例如图像转换）时。

• 最佳实践： 从num_workers=0开始调试，并逐步增加，监控性能。常见的启发式方法建议num_workers = 4 * num_GPU。

• 警告： 工作人员过多会增加 RAM 消耗，并可能导致 CPU 资源竞争，从而降低整个系统的性能。

通过内存锁定加速 CPU-GPU 传输

在DataLoader中将pin_memory=True设置为在 CPU 上分配特殊的“锁定内存”（page-locked memory）。

• 机制： 操作系统无法将此内存交换到磁盘。这允许从 CPU 到 GPU 的异步、直接传输，避免额外的中间复制并减少空闲时间。

• 好处： 加速数据传输到 CUDA 设备，允许 GPU 同时处理和接收数据。

• 不适用的情况： 如果您没有使用 GPU，pin_memory=True不会带来任何好处，只会消耗额外的非可分页 RAM。在 RAM 有限的系统上，可能会对物理内存造成不必要的压力。

实际实施和基准测试

到此为止，我们进入了实验优化 PyTorch 模型训练方法的阶段，比较标准训练循环与高级数据加载技术。

为了展示所讨论方法的有效性，我们考虑了一个包含标准 MNIST 数据集 的前馈神经网络的实验设置。

覆盖的优化技术：

• 标准训练（基线）： PyTorch 中的基本训练周期 (num_workers=0, pin_memory=False)。

• 多工作进程数据加载： 使用多个进程的并行数据加载 (num_workers=N)。

• Pinned Memory + 非阻塞传输： GPU 内存和 CPU-GPU 传输的优化 (pin_memory=True 和 non_blocking=True)。

• 性能分析： 执行时间和最佳实践的对比。

设置测试环境

步骤 1：导入库

第一步是导入所有必要的库并验证硬件配置：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from time import time
import warnings
warnings.filterwarnings('ignore')

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    device = torch.device("cuda")
    print(f"GPU device: {torch.cuda.get_device_name(0)}")
    print(f"GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
else:
    device = torch.device("cpu")
    print("Using CPU")

print(f"Device used for training: {device}")

预期结果：

PyTorch version: 2.8.0+cu126
CUDA available: True
GPU device: NVIDIA GeForce RTX 4090
GPU memory: 25.8 GB
Device used for training: cuda

步骤 2：数据集分析和加载

MNIST 数据集是一个基本基准，包含 70,000 个 28×28 灰度图像。数据归一化对于训练效率至关重要。

让我们定义用于加载数据集的函数：

transform = transforms.Compose()
train_dataset = datasets.MNIST(root='./data',
                               train=True,
                               download=True,
                               transform=transform)

test_dataset = datasets.MNIST(root='./data',
                              train=False,
                              download=True,
                              transform=transform)

步骤 3：为 MNIST 实现一个简单的神经网络

让我们为我们的实验定义一个简单的前馈神经网络：

class SimpleFeedForwardNN(nn.Module):
    def __init__(self):
        super(SimpleFeedForwardNN, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

步骤 4：定义经典训练周期

让我们定义一个可重用的训练函数，它封装了三个关键阶段（正向传递、反向传递和参数更新）：

def train(model,
          device,
          train_loader,
          optimizer,
          criterion,
          epoch,
          non_blocking=False):

    model.train()
    loss_value = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        # Move data on GPU using non blocking parameter
        data = data.to(device, non_blocking=non_blocking)
        target = target.to(device, non_blocking=non_blocking)

        optimizer.zero_grad() # Prepare to perform Backward Pass
        output = model(data) # 1\. Forward Pass
        loss = criterion(output, target)
        loss.backward() # 2\. Backward Pass
        optimizer.step() # 3\. Parameter Update

        loss_value += loss.item()

    print(f'Epoch  {epoch} | Average Loss: {loss_value:.6f}')

分析 1：无优化训练周期（基线）

顺序数据加载的配置 (num_workers=0, pin_memory=False)：

model = SimpleFeedForwardNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Baseline setup: num_workers=0, pin_memory=False
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

start = time()
num_epochs = 5
print("\n==================================================\nEXPERIMENT: Standard Training (Baseline)\n==================================================")
for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch, non_blocking=False)

total_time_baseline = time() - start
print(f"✅ Experiment completed in {total_time_baseline:.2f} seconds")
print(f"⏱️  Average time per epoch: {total_time_baseline / num_epochs:.2f} seconds")

预期结果（基线场景）：

==================================================
EXPERIMENT: Standard Training (Baseline)
==================================================
Epoch  1 | Average Loss: 0.240556
Epoch  2 | Average Loss: 0.101992
Epoch  3 | Average Loss: 0.072099
Epoch  4 | Average Loss: 0.055954
Epoch  5 | Average Loss: 0.048036
✅ Experiment completed in 22.67 seconds
⏱️  Average time per epoch: 4.53 seconds

分析 2：带有工作进程的优化训练循环

我们通过 num_workers=8 引入数据加载的并行性：

model = SimpleFeedForwardNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# DataLoader optimization by using WORKERS
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=8)

start = time()
num_epochs = 5
print("\n==================================================\nEXPERIMENT: Multi-Worker Data Loading (8 workers)\n==================================================")
for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch, non_blocking=False)

total_time_workers = time() - start
print(f"✅ Experiment completed in {total_time_workers:.2f} seconds")
print(f"⏱️  Average time per epoch: {total_time_workers / num_epochs:.2f} seconds")

预期结果（工作进程场景）：

==================================================
EXPERIMENT: Multi-Worker Data Loading (8 workers)
==================================================
Epoch  1 | Average Loss: 0.228919
Epoch  2 | Average Loss: 0.100304
Epoch  3 | Average Loss: 0.071600
Epoch  4 | Average Loss: 0.056160
Epoch  5 | Average Loss: 0.045787
✅ Experiment completed in 9.14 seconds
⏱️  Average time per epoch: 1.83 seconds

分析 3：带有优化和 Pin Memory 的训练循环

我们在 DataLoader 中添加 pin_memory=True，在 train 函数中添加 non_blocking=True 以实现异步传输。

model = SimpleFeedForwardNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Optimization of dataLoader with WORKERS + PIN MEMORY
train_loader = DataLoader(train_dataset,
                          batch_size=64,
                          shuffle=True,
                          pin_memory=True, # Attiva la memoria bloccata
                          num_workers=8)

start = time()
num_epochs = 5
print("\n==================================================\nEXPERIMENT: Pinned Memory + Non-blocking Transfer (8 workers)\n==================================================")
# non_blocking=True for async data transfer 
for epoch in range(1, num_epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch, non_blocking=True)

total_time_optimal = time() - start
print(f"✅ Experiment completed in {total_time_optimal:.2f} seconds")
print(f"⏱️  Average time per epoch: {total_time_optimal / num_epochs:.2f} seconds")

预期结果（所有优化场景）：

==================================================
EXPERIMENT: Pinned Memory + Non-blocking Transfer (8 workers)
==================================================
Epoch  1 | Average Loss: 0.269098
Epoch  2 | Average Loss: 0.123732
Epoch  3 | Average Loss: 0.090587
Epoch  4 | Average Loss: 0.073081
Epoch  5 | Average Loss: 0.062543
✅ Experiment completed in 9.00 seconds
⏱️  Average time per epoch: 1.80 seconds

结果分析和解释

结果展示了数据管道优化对总训练时间的影响。从顺序加载（基线）切换到并行加载（多工作进程）可以将总时间减少超过 50%。添加 带有 Pinned Memory 的 非阻塞 可以提供进一步的小幅但显著的改进。

方法	总时间（秒）	加速比
标准训练 (Baseline)	22.67	基线
多工作进程加载 (8 个工作进程)	9.14	2.48x
优化 (Pinned + 非阻塞)	9.00	2.52x

对结果的反省：

• num_workers 的影响： 引入 8 个工作进程将总训练时间从 **22.67 秒减少到 9.14 秒，实现了 2.48 倍的速度提升。这表明在 基线 情况下，主要瓶颈是数据加载（GPU 的 CPU 饥饿）。

• pin_memory 的影响：添加 pin_memory=True 和 non_blocking=True 进一步将时间减少到 9.00 秒，提供整体性能的提升，最高可达 2.52 倍。这种改进虽然不大，但反映了消除 CPU 锁定内存和 GPU 之间数据传输中的小同步延迟（操作 cudaMemcpyAsync）。

获得的结果并非普遍适用。优化的有效性取决于外部因素：

• 批大小：更大的批大小可以提高 GPU 计算效率，但可能导致内存错误（OOM）。如果发生 I/O 瓶颈，增加批大小可能不会导致训练速度更快。

• 硬件：num_workers 的效率与 CPU 核心数和 I/O 速度（SSD 与 HDD）直接相关。

• 数据集/预处理：应用于数据的转换复杂性影响 CPU 工作量，从而影响 num_workers 的最佳值。

结论

优化神经网络的性能不仅限于选择架构或训练参数。持续监控管道并识别瓶颈（CPU、GPU 或数据传输）可以带来显著的效率提升。

记忆最佳实践

使用如 PyTorch 分析器 等工具进行诊断非常重要。优化 DataLoader 仍然是解决 GPU 空闲问题的最佳起点。

DataLoader 参数	对效率的影响	何时使用
num_workers	并行化预处理和加载，减少 GPU 等待时间。	当分析器指示 CPU 瓶颈时。
pin_memory	加速异步 CPU-GPU 传输。	即，如果你使用 GPU，消除潜在的瓶颈。

超出 DataLoader 的可能未来发展方向

为了进一步加速，你可以探索高级技术：

• 自动混合精度 (AMP)：使用降低精度的（FP16）数据类型来加速计算并将 GPU 内存使用量减半。

• 梯度累积：当 GPU 内存有限时，模拟更大批量大小的技术。

• 专用库：使用类似 NVIDIA DALI 的解决方案，将整个预处理管道移动到 GPU，消除 CPU 瓶颈。

• 硬件特定优化：使用扩展，如 PyTorch 的英特尔扩展，充分利用底层硬件。