Weights & Biases(W&B)是一个用于机器学习实验跟踪与 MLOps 的工具。资料说明,当你需要跟踪 ML 实验并自动记录指标、在实时仪表盘中可视化训练、跨超参数与配置比较运行、使用自动化 sweeps 优化超参数、通过版本与血缘管理模型注册表、在团队工作区中协作 ML 项目,以及跟踪数据集、模型、代码等 Artifacts 的血缘时,可以使用该技能。资料给出的数据为:用户 200,000+ ML 从业者,GitHub Stars 10.5k+,集成 100+。
安装与登录 安装命令为 pip install wandb。登录方式为 wandb login(会创建 API key),或以编程方式设置 API key:export WANDB_API_KEY=your_api_key_here。
快速开始与基本实验跟踪 基本用法是导入 wandb,通过 wandb.init 初始化一次运行,传入 project 与 config(例如 learning_rate、epochs、batch_size、architecture)。在训练循环中调用 wandb.log 记录指标,例如 epoch、train/loss、val/loss、train/accuracy、val/accuracy,最后调用 wandb.finish()。资料还给出 PyTorch 示例:初始化时设置 project 与 config,通过 wandb.config 访问配置,在训练循环中按批次记录 loss、epoch、batch,保存模型后使用 wandb.save("model.pth") 上传到 W&B,最后 wandb.finish()。
核心概念
- 项目与运行:Project 是相关实验的集合,Run 是训练脚本的单次执行。wandb.init 可设置 project、name(可选运行名)、tags(用标签组织)、notes(添加备注)。每次运行有唯一 ID,可通过 run.id 与 run.url 获取。
- 配置跟踪:可自动跟踪超参数,config 中可包含模型架构、训练参数(learning_rate、batch_size、epochs、optimizer)、数据参数(dataset、augmentation)等;训练中可通过 wandb.config.learning_rate、wandb.config.batch_size 访问。
- 指标记录:可记录标量(loss、accuracy)、多个指标(train/loss、train/accuracy、val/loss、val/accuracy、learning_rate、epoch)、自定义 x 轴(wandb.log({"loss": loss}, step=global_step))、媒体(图像、音频、视频,如 wandb.Image)、直方图(wandb.Histogram)、表格(wandb.Table)。
- 模型检查点:可保存包含 epoch、model_state_dict、optimizer_state_dict、loss 的 checkpoint,使用 torch.save 保存后通过 wandb.save('checkpoint.pth') 上传;资料推荐使用 Artifacts:创建 wandb.Artifact('model', type='model'),add_file('checkpoint.pth'),再 wandb.log_artifact(artifact)。
超参数 Sweeps 可自动搜索最优超参数。定义 sweep_config 时,method 可为 'bayes'、'grid' 或 'random';metric 指定 name 与 goal(maximize 或 minimize);parameters 可指定分布与取值,例如 learning_rate 使用 log_uniform(min 1e-5,max 1e-1)、batch_size 使用 values [16, 32, 64, 128]、optimizer 使用 values ['adam', 'sgd', 'rmsprop']、dropout 使用 uniform(min 0.1,max 0.5)。通过 wandb.sweep(sweep_config, project="my-project") 初始化 sweep,定义训练函数(内部 wandb.init() 并通过 wandb.config 访问 sweep 参数),再使用 wandb.agent(sweep_id, function=train, count=50) 运行 50 次试验。资料列出三种策略:grid(穷举)、random(随机搜索)、bayes(贝叶斯优化,推荐)。
Artifacts 与模型注册表 Artifacts 用于跟踪数据集、模型及其他文件的血缘。记录 Artifacts 时可创建 wandb.Artifact(name、type、description、metadata),通过 add_file 或 add_dir 添加文件,再 wandb.log_artifact(artifact)。使用 Artifacts 时,可在 wandb.init 后通过 run.use_artifact('training-dataset:latest') 获取,再 artifact.download() 下载并使用。模型注册表方面,可将模型记录为 artifact(name、type='model'、metadata),add_file('model.pth'),使用 wandb.log_artifact(model_artifact, aliases=['best', 'production']) 记录,并通过 run.link_artifact(model_artifact, 'model-registry/production-models') 链接到模型注册表。
集成示例 资料提供 HuggingFace Transformers 集成:wandb.init 后,在 TrainingArguments 中设置 report_to="wandb"、run_name、logging_steps、save_steps,Trainer 会自动记录到 W&B。PyTorch Lightning 集成:使用 WandbLogger(project="lightning-demo", log_model=True),再传给 Trainer 的 logger。Keras/TensorFlow 集成:wandb.init 后,在 model.fit 的 callbacks 中加入 WandbCallback() 以自动记录指标。
可视化与分析 可记录自定义图表,例如用 matplotlib 生成图后通过 wandb.Image(fig) 记录;也可记录混淆矩阵,使用 wandb.plot.confusion_matrix(probs、y_true、preds、class_names)。资料提到可在 W&B UI 中创建可分享的 Reports,组合运行、图表与文本,支持 Markdown、可嵌入可视化与团队协作。
最佳实践 资料列出若干实践:1. 使用标签与分组组织,例如 tags、group(分组相关运行)、job_type(任务类型)。2. 记录所有相关信息,包括系统指标(gpu/util、gpu/memory、cpu/util)、代码版本(git_commit)、数据划分(data/train_size、data/val_size)。3. 使用描述性名称,例如 name="bert-base-lr0.001-bs32-epoch10",避免 run1 这类通用名称。4. 保存重要 Artifacts,例如最终模型与用于分析的预测表。5. 网络不稳定时使用离线模式:设置 os.environ["WANDB_MODE"] = "offline",之后可用 wandb sync <run_directory> 同步。
团队协作 运行可通过 URL 自动分享(run.url)。团队项目方面,可在 wandb.ai 创建团队账户、添加团队成员、设置项目可见性(私有/公开),并使用团队级 artifacts 与模型注册表。
定价 资料列出:Free 为无限公开项目、100GB 存储;Academic 对学生/研究者免费;Teams 为 $50/seat/month,包含私有项目与无限存储;Enterprise 为定制定价,提供本地部署选项。
资源 文档:https://docs.wandb.ai;GitHub:https://github.com/wandb/wandb(10.5k+ stars);示例:https://github.com/wandb/examples;社区:https://wandb.ai/community;Discord:https://wandb.me/discord。另可参见 references/sweeps.md(超参数优化指南)、references/artifacts.md(数据与模型版本化模式)、references/integrations.md(框架特定示例)。
限制与未提供信息 资料未提供具体的性能基准、支持的全部框架清单、各集成方式的完整限制、离线模式的具体同步细节、企业版功能细节,以及除上述定价外的其他收费信息。