监控和串流机器学习的进度，损失和正样本校验数值

标签

AI/PyTorch

开发/Python/PyTorchLightning

AI/TensorBoard

开发/Python/TensorBoard

开发/Python/wandb

AI/数据集/MNIST

字数

1087 字

阅读时间

5 分钟

要在 PyTorch 监控和可视化训练时候的各项数据，有这么几种方案。

使用 PyTorch 的 Tensorboard 集成

TensorBoard 原本是给 Tensorflow 用于可视化神经网络训练结果的工具，后来大家都挺喜欢的，PyTorch 有一个专门的输出封装来输出数据给 Tensorboard 进行可视化。

主要的能力是：

可视化模型架构
创建交互式的可视化版本
记录训练过程中的各种指标（metrics）
可视化迭代（iteration）过程中的损失（loss）
检查训练数据
绘制 Matplotlib 图表
可视化模型在随机的 minibatch ^[1]上的预测结果

基本的用例是这样的：

python

from torch.utils.tensorboard import SummaryWriter 

# 默认情况下，SummaryWriter 会输出到 `runs` 目录
# 此处我们指定为 runs/fashion_mnist_experiment_1
writer = SummaryWriter('runs/fashion_mnist_experiment_1') 

# 在训练的循环函数中可以这样修改
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # 其他的各种训练步骤...
        running_loss += loss.item()
        # 比如每 1000 个 minibatch 之后记录一次
        if i % 1000 == 999:
			# 记录一个名为 training loss 的指标
            writer.add_scalar('training loss', 
                            running_loss / 1000, 
                            epoch * len(trainloader) + i) 
            running_loss = 0.0

然后另起一个命令行窗口，运行 tensorboard --logdir runs/fashion_mnist_experiment_1 就可以启动一个 Tensorboard 服务来托管指标可视化界面了，接下来在浏览器中打开就可以看到^[2]。

使用 wandb（Weights & Biases）的 PyTorch 集成

wandb（Weights & Biases）是另一个最近比较著名和好用的，可以用于可视化训练进度和对比模型效果的工具。

wandb 可以直接与 PyTorch 混合在一起使用，可以用来记录梯度、度量值、图像等等的各种数据^[3]。

基本的用例是这样的：

python

import wandb 

# 初始化 wandb
wandb.init(project="my_project") 

# 其他正常的配置
model = ...
optimizer = ...

# 如果需要 wandb 自动观察模型变化和梯度的话，可以这样
wandb.watch(model)

# 在训练的循环函数中可以这样修改
for epoch in range(num_epochs):
    for i, data in enumerate(trainloader, 0):
        # 其他的各种训练步骤...

        # 然后在这里记录一下指标
        wandb.log({"loss": loss, "accuracy": accuracy})

接下来登录到 wandb（Weights & Biases）界面上就可以远程监控了^[3:1]！

wandb 确实会很简单，跟随这个 PyTorch | Weights & Biases Documentation 教程很快就能集成到现有的代码里面。但是缺点是这个需要串流到公网的 wandb.ai 实例上，不太方便私有云使用^[4]。

那 wandb 有替代品吗？

Neptune.ai 建议的其他替代品：The Best Weights & Biases Alternatives

如果真的希望离线使用的话，wandb 有 offline（离线模式）可以使用，请参考 General | Weights & Biases Documentation 将 WANDB_MODE 配置为 offline 即可。

Reddit 上讨论的其他替代品：Alternatives to W&B? : MachineLearning，另一家巨头 DagsHub 建议的替代品：ML Experiment Tracking Tools: Comprehensive Comparison | DagsHub，当然 LibHunt 上也可以找到其他的信息：Wandb Alternatives and Reviews。

用 PyTorch Lightning

PyTorch Lightning 是对 PyTorch 的轻量封装（封装了 Trainer 和常用的各种函数和类），和 PyTorch 以及 wandb 不同的是，PyTorch Lightning 默认就会输出好几个 Tensorboard 指标（并且支持在分布式场景下使用），用 TensorBoard 或其他支持的可视化面板就可以^[5]进行可视化了。

基本的用例是这样的：

python

import pytorch_lightning as pl

class LitModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        # training steps here...
        # ...
        self.log("loss", loss)