wandb: 深度学习轻量级可视化工具

wandb全称为Weights & Biases，它是一个用于深度学习实验跟踪和可视化的工具。Weights & Biases提供了一个平台，使机器学习工程师能够更轻松地记录实验的详细信息、跟踪模型性能，并可视化实验结果。通过使用wandb，用户可以更好地理解模型的行为、比较不同实验的效果，并与团队成员共享实验结果。

wandb可以做到：

  • 保存模型训练过程中的超参数
  • 实时可视化训练过程中metric的变化
  • 分析训练过程中系统指标的变化情况，包括CPU/GPU的利用率等
  • 跨平台，与多个深度学习框架的集成，包括Pytorch，Tensorflow等
  • 团队协作，复现历史结果，实验记录的永久保留

wandb主要由四个模块组成，分别是：

  1. dashboard：跟踪和分析实验，提供可视化结果
  2. report：提供实验报告
  3. sweeps：调节和搜索超参数来优化模型
  4. artifacts：数据集和模型版本化

Quick Start

在使用之前，首先得安装相关的依赖库。

1
pip install wandb

之后在命令行中输入如下命令进行登陆，根据提示输入对应的token。token可以在Authorize|Weights & Biases中获取。

1
wandb login

下面我们来模拟一次训练过程，同时在其中使用wandb来记录相关信息：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import wandb
import random
from datetime import datetime

# 初始化wandb来追踪项目project的某次运行run
wandb.init(
    # 设置项目project的名称
    project="my-test-project",

    # 设置该次运行run的相关参数信息
    config={
        "architecture": "CNN",
        "dataset": "CIFAR-100",
        "epochs": 10,
        "batch_size": 2,
        "learning_rate": 0.02
    },
    
    # 设置此次记录的run的名称，例如可以设置为'my-run-1'，这里使用当前时间的字符串形式作为run的名称
    name=datetime.now().strftime('%Y-%m-%d %H:%M:%S')
)

# 模拟训练的过程
epochs = 10
offset = random.random() / 5

for epoch in range(2, epochs):
    # 模拟每次训练得到的准确率和损失
    acc = 1 - 2 ** -epoch - random.random() / epoch - offset
    loss = 2 ** -epoch + random.random() / epoch + offset

    # 调用log，记录相关信息
    wandb.log({"acc": acc, "loss": loss})

# 用于结束此次运行run。这是一个可选操作，但是在notebook中是必须的
wandb.finish()
  • wandb中的层级主要分为两级，第一层级是Project，表示一个项目；第二层级是Run，表示每次程序的运行。
  • 一次运行Run必须从属为某个Project，所有的信息记录都是以Run作为基本单位的。
  • 首先我们需要调用init函数来追踪项目的某次运行，其中需要设置相关的参数。
    • project为对应项目的名称
    • config为项目的相关信息
    • name为记录的run名称，可以手动指定，如果不指定的话会自动生成
  • 在训练过程中，我们会记录过程中的metric，此时需要使用log函数。该函数接收一个dict，其中记录的key-value对，最终会被记录成chart图表的形式。

使用config上传的key-value只会被记录一次，是整个run的描述信息，可以在run -> overview 中查看；

使用log上传的key-value会以序列的形式记录下来，默认可以显示为折线图的形式。

完成上面代码的运行之后，就可以在对应的wandb链接中查看到由acc和loss数据组成的折线图，在overview中可以查看到对应的config信息。如果需要在dashboard中展示，则可以自己手动添加panel，例如使用Scalar chart展示epoch，batch_size等信息。

在某些情况下，例如在调试代码的时候，我们不希望使用wandb进行记录，那么我们可以使用下面的命令来快速启动和关闭wandb。

1
2
3
4
5
6
7
wandb disabled  # 禁用wandb

wandb enabled # 启用wandb

wandb offline # 设置wandb为离线模式

wandb online # 设置wandb为在线模式

Used in Pytorch

基础使用

下面是在pytorch中使用wandb的基础方式。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# import the library
import wandb

# start a new experiment
wandb.init(project="new-sota-model")

# capture a dictionary of hyperparameters with config
wandb.config = {"learning_rate": 0.001, "epochs": 100, "batch_size": 128}

# set up model and data
model, dataloader = get_model(), get_data()

# optional: track gradients
wandb.watch(model)

for batch in dataloader:
    metrics = model.training_step()
    # log metrics inside your training loop to visualize model performance
    wandb.log(metrics)

# optional: save model at the end
model.to_onnx()
wandb.save("model.onnx")

样例记录

在模型训练的过程中，我们每隔一段时间会进行一些验证的工作，来查看当前模型的效果。利用wandb，我们还可以将模型的效果更加直观地展现出来，例如在图像相关的验证过程中，可以调用wandb.Image类来记录相关的结果。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
experiment = wandb.init(project='my-project')
experiment.config.update(
    dict(epochs=epochs, batch_size=batch_size, learning_rate=learning_rate)
)

for epoch in range(1, epochs + 1):
    # Training round
    # ...

    # Evaluation round
    val_score = evaluate(model, val_loader, device, amp)
    scheduler.step(val_score)

    experiment.log({
        'learning rate': optimizer.param_groups[0]['lr'],
        'validation Dice': val_score,
        'images': wandb.Image(images[0].cpu()),
        'masks': {
            'true': wandb.Image(true_masks[0].float().cpu()),
            'pred': wandb.Image(masks_pred.argmax(dim=1)[0].float().cpu()),
        },
        'epoch': epoch,
    })

上面是一个图像分割问题的部分代码。在evaluation round中，我们首先完成基础的验证，然后将其中一张图片的真实值、ground truth和预测结果记录到wandb中。wandb会对图像进行处理和可视化，在wandb中我们就可以很直观地看到对应的预测效果。

总结来说，在模型训练过程中使用wandb，基本有如下步骤：

  1. 初始化wandb，指定对应的project
  2. 使用config，指定对应的超参数，运行信息等
  3. 在训练过程中(every epoch/batch),调用log方法来记录对应的metric
  4. 如果需要，也可以在过程中记录一些预测样例
  5. 在训练完成之后，将模型保存到wandb中

Sweep 参数搜索

上面我们介绍的是wandb中run的基础使用，实际上wandb中还提供了参数搜索和可视化的模块，Sweep。在训练模型的过程中，一个逃不开的步骤就是调参。我们可以手动调参，也可以通过某些方式进行参数搜索，例如随机搜索，grid网格搜索等。而Sweeps集成了参数搜索和可视化功能，可以帮助我们在完成调参的同时得到较好的可视化结果。下面我们将通过一个简单的示例来了解Sweeps的使用方式。

sweep

首先我们需要完成Sweep的定义。实际上，Sweep定义的就是如何进行参数搜索，我们需要提供一系列超参数，参数搜索策略和相应的评估策略。这些配置以嵌套字典的形式进行定义。

1
2
3
4
# 参数搜索策略
sweep_config = {
    'method': 'random'
}

上面首先定义了sweep_config中的method。method指的是参数的搜索策略，可选值包括random,grid和bayes。

1
2
3
4
5
6
# 评估策略
metric = {
    'name': 'loss',
    'goal': 'minimize'
}
sweep_config['metric'] = metric

上面定义了评估策略metric。这里我们需要最小化loss。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
# 等待搜索的参数列表
parameters_dict = {
    'optimizer': {
        'values': ['adam', 'sgd']
    },
    'dropout': {
        'values': [0.3, 0.4, 0.5]
    },
    'learning_rate': {
        # a flat distribution between 0 and 0.1
        'distribution': 'uniform',
        'min': 0,
        'max': 0.1
    },
    'batch_size': {
        # integers between 32 and 256
        # with evenly-distributed logarithms
        'distribution': 'q_log_uniform_values',
        'q': 8,
        'min': 32,
        'max': 256,
    },
    'epochs': {
        'value': 1
    }
}
sweep_config['parameters'] = parameters_dict

上面提供了等待搜索的参数列表parameters。其中每个参数和对应的可选值被组织成key-value的格式，value是一个字典。在字典中可以指定一个待选列表；也可以指定一个候选分布，例如uniform；对于固定的参数，也可以在其中直接指定。

至此，我们就得到了一个嵌套字典sweep_config。

train

接下来，我们需要定义train函数。该函数完成的就是一次模型训练过程，它定义了如何使用某次搜索得到的参数组合。这次的参数组合可以通过config来获取。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import logging
import random

def env_prepare(optimizer, dropout, learning_rate, batch_size):
    logging.info(
        f"using params: [optimizer:{optimizer}, dropout: {dropout}, learning_rate:{learning_rate}, batch_sze:{batch_size}]")


def train_epoch():
    return random.random()


def train(config=None):
    with wandb.init(config=config):
        config = wandb.config
        # 训练环境准备，例如准备dataloader, model, optimizer
        env_prepare(optimizer=config.optimizer,
                    dropout=config.dropout,
                    learning_rate=config.learning_rate,
                    batch_size=config.batch_size)
        # 进行模型迭代训练
        for epoch in range(config.epochs):
            avg_loss = train_epoch()
            wandb.log({'loss': avg_loss, 'epoch': epoch})

train函数的结构基本上是相同的，它接收一个config作为输入，从config中读取得到此次选择的参数，之后利用这些参数完成训练，包括训练环境准备和模型迭代训练。在迭代过程中，调用log进行metric的输出。这里的env_prepare表示接收相关参数，完成训练环境准备；train_epoch表示完成一个epoch的训练，返回对应的loss。

start

上面的工作定义完成之后，就可以开始执行参数搜索了，代码如下：

1
2
sweep_id = wandb.sweep(sweep_config, project="my-test-project")
wandb.agent(sweep_id, train, count=10)

开始运行代码，程序就会开始执行，开始进行参数搜索。最终运行完毕之后，可以在wandb相关链接中找到结果。wandb提供了优秀的可视化结果，来帮助我们进行超参数的挑选。

参考文章

  1. PyTorch | Weights & Biases Documentation
  2. Wandb不可缺少的机器学习分析工具-51CTO.COM
  3. Tune hyperparameters | Weights & Biases Documentation
深度学习 > Coding
#深度学习
wandb: 深度学习轻量级可视化工具
http://example.com/2023/11/20/wandb-深度学习轻量级可视化工具/
作者
EverNorif
发布于
2023年11月20日
许可协议