跳转至

tracking.py — 实验追踪

文件路径: verl/utils/tracking.py

文件概述

tracking.py 提供多后端实验追踪支持,统一封装了 WandB、MLflow、TensorBoard、AIM、SwanLab 等常见的实验管理工具。训练过程中的 loss、reward、学习率等指标通过统一接口记录。

核心类

Tracking — 统一追踪接口

class Tracking:
    def __init__(self, project_name, experiment_name, default_backend, config=None):
        """初始化追踪器,根据配置选择后端"""
        self.backends = []
        if "wandb" in default_backend:
            self.backends.append(WandbTracker(...))
        if "mlflow" in default_backend:
            self.backends.append(MLflowTracker(...))
        if "tensorboard" in default_backend:
            self.backends.append(TensorBoardTracker(...))
        # ...

    def log(self, data: dict, step: int):
        """记录指标到所有后端"""
        for backend in self.backends:
            backend.log(data, step)

    def finish(self):
        """结束追踪"""
        for backend in self.backends:
            backend.finish()

可以同时使用多个后端(如同时写 WandB 和 TensorBoard)。

支持的后端

后端 说明
WandB 最流行的实验管理平台
MLflow 开源 ML 生命周期管理
TensorBoard PyTorch/TensorFlow 标配
AIM 高性能实验追踪
SwanLab 国产实验追踪

与其他模块的关系

  • 被 PPO trainer 在每个训练 step 调用
  • 依赖 import_utils.py 检查各后端是否安装

小结

tracking.py 通过统一的 Tracking 类屏蔽了不同实验追踪后端的差异,用户只需在配置中指定后端名称即可。