训练可视化工具哪款是你的菜？MMCV一行代码随你挑

在深度学习中可视化模型的训练过程有助于我们分析模型的状态。可视化训练过程的库很多，我们将一些常用的库集成到 MMCV 中方便用户使用。在 MMCV 中使用这些库只需简单配置。在本文中将介绍这些库以及它们在 MMCV 中的使用方法。

TensorBoard

https://pytorch.org/docs/stable/tensorboard.html

介绍

TensorBoard 最初是随 TensorFlow 提出的一款可视化工具包，其便捷性和完善的记录功能使它得到了广泛应用，并扩展到 PyTorch 等多种深度学习框架。TensorBoard 支持记录多种数据类型：

- 指标和损失

- 超参数和模型 config

- 图片数据（可视化权重、张量、多个图像）

- 模型图

- Embedding Projector（在低维空间可视化高维数据）

作为一款可视化工具，TensorBoard 可以基本满足深度学习实验的记录和可视化需求，它提供了本地端口展示可视化结果，使用起来十分便捷。但也正因为 TensorBoard 的主要功能是可视化，如果涉及到纷繁复杂的实验管理以及机器学习生命周期的记录，我们可能需要借助于其他工具。

使用

在 OpenMMLab codebase 中 使用 TensorBoard 只需一行配置，举 MMClassification 为例：

- 安装 MMClassification

- 安装 TensorBoard

OpenMMLab codebase 网址：

https://github.com/open-mmlab

MMClassification 网址：

https://github.com/open-mmlab/mmclassification

pip install tensorboard

- 修改 MMClassification 中的 log_config 字段

log_config = dict( 
    interval=100, 
    hooks=[ 
        dict(type='TextLoggerHook'), 
        dict(type='TensorboardLoggerHook') 
    ])

- 运行命令

python tools/train.py configs/resnet/resnet18_b16x8_cifar10.py 
# 另起一个终端 
tensorboard --logdir work_dirs/resnet18_b16x8_cifar10

- 打开 chrome 浏览器 输入 http://localhost:6006/

-如果想分享 tensorboard 日志，可以在上面的 tensorboard 命令中加入 dev 参数）

Neptune

https://neptune.ai/

介绍

Neptune 是一个集实验记录、数据存储、可视化、模型注册等多种功能于一体的机器学习实验管理工具，用户可以在网页端 UI 轻松地查看所有的记录与可视化结果。Neptune 支持记录的数据类型包括但不限于：

- 指标和损失

- 超参数和模型 config

- 模型 checkpoints

- Git 信息

- 数据版本管理

- 硬件消耗

- 文件

- 控制台日志

- 图片数据（图片文件、Matplotlib figure、PIL image、Numpy array、Tensor）

- 交互式可视化（自动将 Matplotlib figure 转为交互式，同时支持其他格式如 html 文件、Altair chart）

相较于 TensorBoard，Neptune 支持记录更多种类的数据，并且提供了用户友好的 UI，使用户可以灵活地调整可视化界面。Neptune 还提供了 TensorBoard 接口，可以很方便地把 TensorBoard logs 转换为 Neptune experiments。

使用

- 安装 neptune

pip install neptune-client

- 注册 neptune 账号并设置 NEPTUNE_API_TOKEN

- 修改 log_config 字段，其中 YOUR_WORKSPACE 是账号名，YOUR_PROJECT 是项目名

NEPTUNE_API_TOKEN 网址：

https://docs.neptune.ai/getting-started/installation#authentication-neptune-api-token

log_config = dict( 
    interval=100, 
    hooks=[ 
        dict(type='TextLoggerHook'), 
        dict(type='NeptuneLoggerHook', init_kwargs=dict(project='<YOUR_WORKSPACE/YOUR_PROJECT>')) 
    ])  

- 运行命令

python tools/train.py configs/resnet/resnet18_b16x8_cifar10.py

- 效果图

Weights & Biases

https://wandb.ai/site

介绍

Weights & Biases（简称 WandB）也是一个综合了多种功能的机器学习实验管理工具，其基本功能和Neptune 类似，也有一些细微差别。WandB 支持记录的数据类型包括但不限于：

- 指标和损失

- 超参数和模型 config

- 硬件消耗

- Git 信息

- 控制台日志

- 数据集

- 模型参数和梯度

WandB 的一大优势是提供了 PyTorch / TensorFlow / Keras 的 warpper，通过 wrapper 可以自动地记录模型中每一层的 weights / biases / gradients 。另外，WandB 的 reports 功能也非常实用，它可以将可视化结果以及用户的笔记整理成 report，便于归档整理和团队协作。WandB 也提供了 TensorBoard 的集成功能。

使用

- 安装 wandb

pip install wandb

- 注册 wandb 账号并获取 API key

- 终端输入 wandb login 回车并输入 API key

- 修改 log_config 字段

log_config = dict( 
    interval=100, 
    hooks=[ 
        dict(type='TextLoggerHook'), 
        dict(type='WandbLoggerHook', init_kwargs=dict(project='Your-project')), 
    ])  

- 运行命令

python tools/train.py configs/resnet/resnet18_b16x8_cifar10.py  

- 效果图

MLflow

https://mlflow.org/

介绍

MLflow 是一个用于记录机器学习生命周期的开源工具，实验记录和可视化只是其中一个基础功能，因此它的可视化功能不如 Neptune 和 WandB 那样丰富灵活。MLflow 支持记录的数据类型有：

- 指标和损失

- 超参数和模型config

- Git信息

- Artifacts（图片、模型、数据等）

MLflow 只能以 artifacts 的形式记录图片，没有交互式功能，因此很难从图片中直接获取实验数据。MLflow 也不适用于大型实验，过多的实验可能导致 UI 滞后。然而，MLflow 的主要优势在于机器学习生命周期的完整记录，包括实验可复现性的实现、模型注册、模型和数据的版本管理等。

使用

- 安装 mlflow

pip install mlflow  

- 修改 log_config 字段，其中 exp_name 是项目名

log_config = dict( 
    interval=100, 
    hooks=[ 
        dict(type='TextLoggerHook'), 
        dict(type='MlflowLoggerHook', exp_name='your-project'), 
    ])  

- 运行命令

python tools/train.py configs/resnet/resnet18_b16x8_cifar10.py 
# 另起一个终端 
mlflow ui

- 打开 chrome 浏览器 输入 http://localhost:5000/

- 效果图

Dvclive

https://dvc.org/doc/dvclive

介绍

Dvclive 是一个用于追踪机器学习模型训练过程的开源 Python 库。Dvclive 与 DVC 无缝集成，它生成的 logs 可以作为 dvc plots；另一方面，Dvclive 也可以独立于 DVC 工作，TSV 格式的 log 文件可以通过其他可视化工具转换为图表。Dvclive 主要用于记录指标和损失。在与 DVC 集成后，可以通过plots 功能将 log 可视化，用户可以选择折线图等多种绘图模板，也可以自定义模板。

作为可视化工具，Dvclive 的功能是有限的。因此在大多数情况下，它与 DVC 集成使用，用于记录实验过程并自动生成 training summaries。

使用

- 安装

pip install dvc  

- 修改 log_config 字段

log_config = dict( 
    interval=100, 
    hooks=[ 
        dict(type='TextLoggerHook'), 
        dict(type='DvcliveLoggerHook', path='dvclive'), 
    ])

- 运行命令

dvc init 
python tools/train.py configs/resnet/resnet18_b16x8_cifar10.py 
# 另起一个终端 
dvc plots show dvclive/train/loss.tsv

- 效果图，另外也可以搭配 DVC Studio 使用

DVC Studio 网址：

https://studio.iterative.ai/team/Iterative/views/example-get-started-zde16i6c4g

总结

最后对以上工具做一个简单的总结：

如果你偏好在本地使用，想要一款能满足基本的实验记录需求、且上手容易的，推荐使用 TensorBoard；

如果对可视化和实验记录有较高要求，推荐使用 Neptune 和 WandB 。两者的实验记录功能都非常强大，并且都支持多人协作，适合大型项目的管理。

如果你是开源项目的爱好者，或者你希望记录完整的 ML 实验全过程，MLflow是一个不错的选择。

如果你只需要记录基本的实验数据，但是对于大规模数据的版本管理有比较高的需求，那么推荐你使用 DVC。