使用 Swanlab 和 WandB 追踪实验数据

Swanlab 和 WandB 是两个强大的 AI 实验跟踪工具，支持多种框架，能够高效记录和分析实验数据，无论是在本地开发环境还是云端协作场景，都可以帮助开发者提升工作效率。

本文旨在帮助开发者掌握如何在 AIStudio 开发机和任务上正确配置和使用 Swanlab 和 WandB 第三方工具，以满足实验监控和可视化的需求。

创建实验环境

以下教程主要以 AIStudio 开发机为例。进入创建页面后，请根据页面提示，完成开发机的相应配置。请开启 SSH 登录。

点击下方链接，直接访问智算云控制台的创建开发机页面。

创建开发机

SwanLab

安装 SwanLab

pip install swanlab

获取 SwanLab API key

如果要同步实验数据到云端（https://swanlab.cn/），必须注册 SwanLab 账号并获取 SwanLab API key。

推荐使用环境变量登录 SwanLab：

export SWANLAB_API_KEY=TjGzbMlAx5dW3QeFUgj9P

请替换为您自己的 API key。

NOTE

SwanLab 也支持使用 Shell 命令或在 Python 代码中登录，详见 SwanLab 登录账号。

IMPORTANT

如果您不使用 SwanLab 云端，仅使用离线看板功能，那么您不需要获取 API key。

用 SwanLab 记录实验数据

1. 使用 swanlab.init() 初始化一个 SwanLab 实验。

   import swanlab
   
   run = swanlab.init(
       project = "my-project", # 项目名称
       # experiment_name = "test1", # 本轮实验名称，如果不指定则默认取 "exp" 拼接时间戳
       # description = "我的第一个人工智能实验", # 实验介绍
   )

   NOTE

   更多 swanlab.init() 参数详见 Swanlab

2. 使用 swanlab.log() 收集指标名和数据。

   import swanlab
   import random
   
   # 初始化SwanLab
   run = swanlab.init(
       # 设置项目
       project="my-project",
       # 跟踪超参数与实验元数据
       config={
           "learning_rate": 0.01,
           "epochs": 10,
       },
   )
   
   print(f"学习率为{run.config.learning_rate}")
   
   offset = random.random() / 5
   
   # 模拟训练过程
   for epoch in range(2, run.config.epochs):
       acc = 1 - 2**-epoch - random.random() / epoch - offset
       loss = 2**-epoch + random.random() / epoch + offset
       print(f"epoch={epoch}, accuracy={acc}, loss={loss}")
       # 记录指标
       swanlab.log({"accuracy": acc, "loss": loss})

   在训练循环中，将指标名和数据组成一个键值对字典，传递给 swanlab.log({指标名: 指标}) 完成 1 次指标的记录。指标值可以是 int/float/BaseType 三种类型之一。

   关于 swanlab.log() 的用法，建议访问 SwanLab 官方文档记录实验指标。

云端查看实验数据

访问 https://swanlab.cn/workspace 查看记录情况。

本地查看实验数据

实验时，SwanLab 会创建一个本地目录（默认名称为“swanlog”），所有日志和文件都保存在其中，并异步传输到 SwanLab 服务器。SwanLab 还支持在不联网的情况下跟踪实验，以及访问实验记录。

如果您只想将训练数据储存在本地，而不上传到服务器中，您只需更改 swanlab.init() 中的 mode="local" 参数即可，您也可以通过设置参数 logdir 来指定日志的储存路径。

1. 设置 mode="local"（本地模式）。

   run = swanlab.init(
       project="local-test",
       logdir="/root/swanlog", # 指定日志文件储存路径，默认为 swanlog
       mode="local", # 设置为本地模式
       config={
           "learning_rate": 0.01,
           "epochs": 10,
       },
   )

   NOTE

   本地查看实验数据仅支持离线看板，仅支持查看实验数据，不支持工作空间、项目等云端特性。

2. 启动 SwanLab 查看本地日志文件

   在命令行中输入 swanlab watch /root/swanlog，请替换为您的指定文件路径。

3. 本地端口转发

   为了在本地的浏览器中打开离线看板，请在本地计算机命令行终端设置端口转发。

   格式

   ssh -N -f -L 8080:127.0.0.1:5092 -p <开发机 SSH 端口> root@<开发机 SSH IP>

   示例，请务必替换为您的开发机对应配置。

   ssh -N -f -L 8080:127.0.0.1:5092 -p 43282 root@36.213.48.208

   NOTE

   详见 SSH 端口转发。

WandB

安装 WandB

pip install wandb

获取 WandB API key

在 WandB 官网注册和登录账户，并获取您的 API_key。

输入命令，并根据提示输入您的 API_key，之后无需再次登录。

wandb login

NOTE

WandB 也支持在 Python 代码中登录 wandb.login(api_key = "API_key")。

用 WandB 记录实验数据

1. 使用 wandb.init() 初始化实验。

   run = wandb.init(
       # 设置项目名称
       project="my-project",
       # 跟踪超参数和实验元数据
       config={
           "learning_rate": 0.01,
           "epochs": 10,
       },
   )

   NOTE

   在wandb.init()中，您还可以添加其余参数，详见WandB官网。

2. 使用 wandb.log() 记录实验指标。

   import wandb
   import random
   
   wandb.login()
   
   epochs = 10
   lr = 0.01
   
   run = wandb.init(
       # 设置项目
       project="my-project",
       # 跟踪超参数和实验元数据
       config={
           "learning_rate": 0.01,
           "epochs": 10,
       },
   )
   
   offset = random.random() / 5
   print(f"lr: {lr}")
   
   # 模拟训练过程
   for epoch in range(2, epochs):
       acc = 1 - 2**-epoch - random.random() / epoch - offset
       loss = 2**-epoch + random.random() / epoch + offset
       print(f"epoch={epoch}, accuracy={acc}, loss={loss}")
       # 记录指标
       wandb.log({"accuracy": acc, "loss": loss})

云端查看实验数据

访问 https://wandb.ai/home 查看记录情况。

在 AIStudio 任务中使用

除了在开发机中追踪实验数据，平台还支持在 AIStudio 任务中使用 SwanLab 云端版和 Tensorboard 进行可视化监测。

SwanLab 云端版

推荐使用任务的环境变量功能传入 SwanLab 云端版的 API key（SWANLAB_API_KEY）。训练开始后，请访问 swanlab.cn 查看训练数据。

WandB

WandB 的使用方法与 SwanLab 云端版类似。

参考资料

• SwanLab and Transformers: Power Up Your NLP Experiments