PyTorch 中实现模型训练看板实时监控训练过程中的关键指标

小赖同学啊2025-03-06 10:26

在 PyTorch 中实现模型训练看板（Dashboard）可以帮助你实时监控训练过程中的关键指标（如损失、准确率、学习率等）。常用的工具包括 TensorBoard 、Weights & Biases (W&B) 和 Matplotlib。以下是使用这些工具实现训练看板的详细方法。

1. 使用 TensorBoard

TensorBoard 是 PyTorch 官方推荐的训练可视化工具，支持实时监控训练过程中的各种指标。

安装 TensorBoard

bash 复制代码 
pip install tensorboard

在 PyTorch 中使用 TensorBoard

以下是一个简单的示例，展示如何在训练过程中记录损失和准确率。

python 复制代码 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return self.fc(x.view(x.size(0), -1))

# 初始化模型、损失函数和优化器
model = SimpleModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 加载数据集
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 初始化 TensorBoard
writer = SummaryWriter('runs/experiment_1')

# 训练循环
for epoch in range(5):
    running_loss = 0.0
    correct = 0
    total = 0

    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        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()

        if i % 100 == 99:  # 每 100 个 batch 记录一次
            avg_loss = running_loss / 100
            accuracy = 100 * correct / total
            print(f'Epoch {epoch + 1}, Batch {i + 1}: Loss={avg_loss:.4f}, Accuracy={accuracy:.2f}%')

            # 将损失和准确率写入 TensorBoard
            writer.add_scalar('Training Loss', avg_loss, epoch * len(train_loader) + i)
            writer.add_scalar('Training Accuracy', accuracy, epoch * len(train_loader) + i)

            running_loss = 0.0
            correct = 0
            total = 0

# 关闭 TensorBoard
writer.close()

启动 TensorBoard

在终端运行以下命令启动 TensorBoard：

bash 复制代码 
tensorboard --logdir=runs

然后在浏览器中打开 http://localhost:6006 查看训练看板。

2. 使用 Weights & Biases (W&B)

Weights & Biases 是一个强大的实验跟踪工具，支持实时监控训练过程，并且可以与其他团队成员共享结果。

安装 W&B

bash 复制代码 
pip install wandb

在 PyTorch 中使用 W&B

以下是一个简单的示例：

python 复制代码 
import wandb
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 初始化 W&B
wandb.init(project="pytorch-dashboard-example", config={
    "learning_rate": 0.01,
    "epochs": 5,
    "batch_size": 64
})

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return self.fc(x.view(x.size(0), -1))

# 初始化模型、损失函数和优化器
model = SimpleModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=wandb.config.learning_rate)

# 加载数据集
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=wandb.config.batch_size, shuffle=True)

# 训练循环
for epoch in range(wandb.config.epochs):
    running_loss = 0.0
    correct = 0
    total = 0

    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        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()

        if i % 100 == 99:  # 每 100 个 batch 记录一次
            avg_loss = running_loss / 100
            accuracy = 100 * correct / total
            print(f'Epoch {epoch + 1}, Batch {i + 1}: Loss={avg_loss:.4f}, Accuracy={accuracy:.2f}%')

            # 将指标记录到 W&B
            wandb.log({
                "Epoch": epoch + 1,
                "Batch": i + 1,
                "Training Loss": avg_loss,
                "Training Accuracy": accuracy
            })

            running_loss = 0.0
            correct = 0
            total = 0

# 完成训练
wandb.finish()

查看 W&B 看板

运行代码后，可以在 W&B 的网页端查看实时训练看板。

3. 使用 Matplotlib

如果你需要一个简单的本地可视化工具，可以使用 Matplotlib 绘制训练过程中的指标。

安装 Matplotlib

bash 复制代码 
pip install matplotlib

在 PyTorch 中使用 Matplotlib

以下是一个简单的示例：

python 复制代码 
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(28 * 28, 10)

    def forward(self, x):
        return self.fc(x.view(x.size(0), -1))

# 初始化模型、损失函数和优化器
model = SimpleModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 加载数据集
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

# 记录训练过程中的指标
losses = []
accuracies = []

# 训练循环
for epoch in range(5):
    running_loss = 0.0
    correct = 0
    total = 0

    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        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()

        if i % 100 == 99:  # 每 100 个 batch 记录一次
            avg_loss = running_loss / 100
            accuracy = 100 * correct / total
            print(f'Epoch {epoch + 1}, Batch {i + 1}: Loss={avg_loss:.4f}, Accuracy={accuracy:.2f}%')

            losses.append(avg_loss)
            accuracies.append(accuracy)

            running_loss = 0.0
            correct = 0
            total = 0

# 绘制损失和准确率曲线
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(losses, label='Training Loss')
plt.xlabel('Batch')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(accuracies, label='Training Accuracy')
plt.xlabel('Batch')
plt.ylabel('Accuracy')
plt.legend()

plt.show()

总结

  • TensorBoard：适合本地实时监控，功能强大，支持多种可视化。
  • Weights & Biases：适合团队协作和实验跟踪，支持云端存储和共享。
  • Matplotlib：适合简单的本地可视化，无需额外依赖。

根据你的需求选择合适的工具来实现模型训练看板。

相关推荐
yourkin66639 分钟前
什么是神经网络？
人工智能·深度学习·神经网络
嘀咕博客1 小时前
Frames：Runway推出的AI图像生成模型，提供前所未有的风格控制和视觉一致性
人工智能·ai工具
isNotNullX2 小时前
ETL详解：从核心流程到典型应用场景
大数据·数据仓库·人工智能·架构·etl
薰衣草23332 小时前
力扣——位运算
python·算法·leetcode
科技峰行者2 小时前
通义万相2.5系列模型发布，可生成音画同步视频
人工智能·阿里云·ai·大模型·agi
两只程序猿2 小时前
数据可视化 | Violin Plot小提琴图Python实现 数据分布密度可视化科研图表
开发语言·python·信息可视化
Vizio＜2 小时前
《面向物理交互任务的触觉传感阵列仿真》2020AIM论文解读
论文阅读·人工智能·机器人·机器人触觉
尤超宇3 小时前
基于卷积神经网络的 CIFAR-10 图像分类实验报告
人工智能·分类·cnn
alex1003 小时前
BeaverTails数据集：大模型安全对齐的关键资源与实战应用
人工智能·算法·安全
大模型真好玩3 小时前
架构大突破! DeepSeek-V3.2发布，五分钟速通DeepSeek-V3.2核心特性
人工智能·python·deepseek
热门推荐
01OpenSpeedy简介02在国行 macOS 下用 DeepSeek 补齐 Xcode 26 的 AI 能力：问题、原因与 mitmproxy 解决方案（含可用脚本与安装教程）03UV安装并设置国内源04GitHub 镜像站点05UV 工具安装与国内镜像源配置指南06KGG转MP3工具|非KGM文件|解密音频07阿里最新开源Wan2.2-Animate-14B 本地部署教程：统一双模态框架，MoE架构赋能电影级角色动画与替换08VsCode远程Copilot无法使用Claude Agent问题09Spec-Kit 使用指南1046个Nano-banana 精选提示词，持续更新中