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 类屏蔽了不同实验追踪后端的差异,用户只需在配置中指定后端名称即可。