在 PyTorch 训练中使用 `tqdm` 显示进度条

news2024/11/27 20:01:17

在 PyTorch 训练中使用 `tqdm` 显示进度条

在深度学习的训练过程中，实时查看训练进度是非常重要的，它可以帮助我们更好地理解训练的效率，并及时调整模型或优化参数。使用 tqdm 库来为训练过程添加进度条是一个非常有效的方式，本文将介绍如何在 PyTorch 中结合 tqdm 来动态显示训练进度。

1. 安装 `tqdm` 库

首先，如果你还没有安装 tqdm，可以通过 pip 命令进行安装：

pip install tqdm

tqdm 是一个非常轻量级的 Python 库，它可以快速地为循环加上进度条，并提供非常友好的终端显示效果。

2. 如何在训练循环中使用 `tqdm`

2.1 包装 `train_loader` 和显示进度条

在训练过程中，最常见的任务是通过数据加载器（train_loader）来批量读取训练数据并进行前向传播、反向传播、优化等操作。为了显示训练的进度条，我们可以使用 tqdm 来包装 train_loader。这样，我们就能在每次读取数据时自动更新进度条。

示例代码：

from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim

# 假设模型、损失函数、优化器等已经定义好
model = MyModel().to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

# 假设 train_loader 已经定义好
train_loader = ...

num_epochs = 10

# 训练循环
for epoch in range(num_epochs):
    model.train()  # 设置模型为训练模式
    running_loss = 0.0
    correct = 0
    total = 0

    # 使用 tqdm 包装 train_loader，自动显示进度条
    for batch_idx, (audio, labels) in enumerate(tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}", ncols=100)):
        audio = audio.to(device)
        labels = labels.to(device)

        # 前向传播
        optimizer.zero_grad()
        outputs = model(audio)

        # 计算损失
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()
        optimizer.step()

        # 更新统计信息
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    # 输出每个 epoch 的总结信息
    print(f"\nEpoch {epoch+1} complete. Loss: {running_loss/len(train_loader):.4f}, Accuracy: {100 * correct / total:.2f}%")

2.2 进度条的显示效果

当你使用 tqdm 包装 train_loader 时，进度条会自动显示。每个 epoch 内，tqdm 会根据 train_loader 的批次数量动态更新进度条。显示的内容通常包括：

当前 epoch 和 batch 的编号
已完成的批次数量
当前 batch 的损失值
进度百分比
剩余时间估计

进度条示例：

Epoch 1/10: 100%|█████████████████████████████████████████████████████████████████████| 2043/2043 [03:30<00:00,  9.70it/s]

进度条字段说明：

Epoch 1/10: 表示当前是第 1 个 epoch，总共有 10 个 epoch。
100%: 表示该 epoch 的训练已经完成。
█████████████████████████████████████████████████████████████████████: 代表进度条，显示当前训练的进度。
2043/2043: 当前 epoch 中处理的 batch 数，表示已处理 2043 个 batch，总共有 2043 个。
03:30<00:00: 训练已经花费的时间（03:30）和预计剩余时间（<00:00）。
9.70it/s: 表示每秒处理的批次数（即训练速度）。