DL——WandB使用笔记

整体介绍

  • WandB（Weights & Biases）是一个用于机器学习/深度学习项目的实验跟踪、可视化和模型管理的工具
    • 实验跟踪 ：超参数、训练 Loss、评估 Loss 和其他自定义指标
      • 注：可实时更新图表和指标
    • 可视化 ：生成如 Loss 曲线和 AUC 曲线等，同时可以可视化权重分布或特征
      • 注：还可以使用 plt 等自定义曲线并上报在 WandB 上生成
    • 模型管理 ：保存和共享模型的所有相关文件，包括权重文件、配置文件和环境设置，方便重新训练或共享
      • 注：支持模型版本控制，可对模型版本进行跟踪，比较不同模型的性能
    • 团队协作 ：团队成员可以查看实验结果和进度
    • 资源监控 ：可以监控GPU使用率、内存占用等系统资源
      • 注：原生的 Torch 等没有这种 Profile 功能
  • WandB 具有强大的兼容性 ：
    • 适配 Jupyter、TensorFlow、PyTorch、Keras、Scikit-learn、LightGBM、XGBoost 等多种框架

基本使用流程介绍

  • 首先在wandb官网注册账号并获取API key，使用命令wandb login进行登录
  • 使用wandb.init()启动新项目，通过wandb.config跟踪超参数
  • 在训练过程中，使用wandb.log()记录指标，使用wandb.summary()添加最终结果
  • 最后使用wandb.finish()结束实验
  • 基本代码示例见附录

WandB API wandb.log()

整体定义

  • wandb.log() API用于记录和可视化实验指标、模型参数等数据，其基本定义为： 
    1
    wandb.log(data, step=None, commit=True, sync=False)

data参数（必需）

  • 字典（dict）类型，指定要记录的数据，键为字符串（指标名称），值可以是多种类型：
    • 原生类型 ：如 int, float等，用于记录损失、准确率等随训练迭代变化的值
    • 图像 ：wandb.Image() 对象，上传图片后可以根据 step 拉杆选择指定 step 的图片查看
    • 表格 ：wandb.Table() 对象
    • 直方图 ：wandb.Histogram() 对象
    • 其他 WandB 对象 ：如音频、视频、文本等均可
  • 代码示例： 
    1
    2
    3
    4
    5
    # 记录标量值
    wandb.log({"loss": 0.235, "accuracy": 0.92})
    # 记录图像
    img = wandb.Image("path/to/image.jpg", caption="Sample Image")
    wandb.log({"examples": img})

step参数（可选）重点参数

  • 类型为整数（int），用于手动指定当前记录的全局步数（如迭代次数、轮数）。若不指定，WandB 会自动递增
  • 注意点1 ：若 step 参数缺失，则每次调用都会自动生成新的 step（step 自动加 1）
    • 常见问题：一个 step 中，多次调用 wandb.log() 但不明确传入 step 参数，此时模型会错误的认为每次调用 wandb.log() 都是一个新的 step
  • 注意点2 ：新增的 step 必须大于等于已有 step 的最大值
    • 若开始调用时使用了 step=400 的值，则以后再调用 step=300 时是不会记录数据的，如果是依次调用 step=1，step=10，step=8，则仅保留step=1 和 step=10 的数据（step=8会消失）
    • 新增的最小单元是：指标 x step；也就是说，无论是指标还是 step 变化，都算是新增
    • 任意指标新增 step 时，都不能小于当前所有指标的最大 step，但是如果这个指标的这个 step 已经存在，则是可以被修改的
  • 注意点3 ：多次上传同一 step 会覆盖之前的值
    • 重复上传相同 step ，且仅会覆盖之前的值（仅针对相同指标）
  • 使用建议：在需要管理多个指标时，建议加上 step 参数

commit参数（可选）

  • 布尔值（bool），默认 True，用于决定是否立即将当前日志提交到 WandB 服务器
    • 若为 False，则缓存日志，直到下一次 commit=True 或手动调用 wandb.log()

sync参数（可选）

  • 布尔值（bool），默认 False，用于决定是否强制立即同步所有缓存的日志到服务器（即使 commit=False）

使用技巧笔记

  • 使用 wandb.init(mode="offline") 可在本地保存日志，后续手动同步
  • WandB 免费版有每月 10GB 存储限制，大量图片可能快速耗尽配额，建议上传前压缩图片（如 wandb.Image(img, mode="L") 转为灰度图）
  • WandB 可以部署到本地 Docker 上以保证数据安全

附录：WandB会存储哪些信息？

config.yaml

  • 存储训练任务的配置参数（超参数），例如：
    • 模型架构参数（如层数、神经元数量）
    • 训练参数（学习率、批次大小、迭代次数）
    • 数据路径、预处理设置
    • 优化器配置（如SGD、Adam的参数）
    • Python 版本信息
  • 用于记录实验的完整配置，便于复现和对比不同实验
  • 注：wandb.config 中配置的参数都会在 config.yaml 中出现

output.log

  • 包含训练过程中的标准输出日志
  • 注意：初始化了 WandB 后，标准输出会自动记录下来并上传

requirements.txt

  • 记录项目依赖的 Python 库及其版本，例如：

    1
    2
    3
    torch==2.0.0
    numpy==1.24.0
    pandas==2.0.3

  • 用于确保实验环境可复现，其他用户可通过该文件安装相同依赖

  • 与本地 pip list 结果完全一致

wandb-metadata.json

  • 包含 WandB 自动生成的元数据，例如：
    • 实验启动时间、持续时长
    • 运行环境信息（操作系统、Python版本）
    • 硬件信息（CPU/GPU型号、内存）
    • 代码仓库信息（如Git commit哈希）
  • 用于记录实验的环境上下文，便于追溯和管理
  • 示例： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    {
        "os": "macOS-14.4.1-x86_64-i386-64bit",
        "python": "CPython 3.10.0",
        "startedAt": "2024-05-27T06:36:31.346712Z",
        "program": "/Users/xxx/wandb_demo.py",
        "email": "xxx@gmail.com",
        "root": "/Users/xxx/xxx/",
        "host": "MBP-xxx.local",
        "executable": "/Users/xxx/envs/xxx/bin/python"
    }

wandb-summary.json

  • 存储训练过程中的关键指标 Summary
  • 用于快速查看实验结果，无需解析完整日志

附录：代码示例

  • 一个简单的使用 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
    import random
    import wandb
    import time

    # Start a new wandb run to track this script.
    run = wandb.init(
        # entity="xxx", 不指定该项时默认走登录账号默认的组织
        project="wandb-demo-project", # Set the wandb project where this run will be logged.
        name="jiahong_sleep_log", # exp name，如果不指定则默认生成一个
        config={ # Track hyperparameters and run metadata. 会存储到 config.yaml中
            "learning_rate": 0.04,
            "dataset": "gsm8k",
            "epochs": 1,
        },
    )

    def model_step_simulation(step):
        acc = 1 - 2 ** -step - random.random() / step - offset
        loss = 2 ** -step + random.random() / step + offset
        return acc, loss

    # Simulate training.
    steps = 50
    offset = random.random() / 5
    for step in range(2, steps):
        acc, loss = model_step_simulation(step) # 真实场景这里会进行模型训练
        time.sleep(2)

        # Log metrics to wandb.
        wandb.log({"acc": acc, "loss": loss}) # 上报 scalar 图，这里会生成两个图
        print(f"step: {step}, loss: {loss}") # 打印到标准输出，可在 output.log 中查看

    # Finish the run and upload any remaining data.
    wandb.finish()