记录到 Weights & Biases
本深入探讨将指导您完成如何在 torchtune 中设置记录到 Weights & Biases (W&B)。
如何开始使用 W&B
如何使用
WandBLogger
如何将配置、指标和模型检查点记录到 W&B
torchtune 支持将您的训练运行记录到 Weights & Biases。下图显示了来自 torchtune 微调运行的 W&B 工作区示例。
注意
您需要安装 wandb
包才能使用此功能。您可以通过 pip 安装它
pip install wandb
然后,您需要使用 W&B CLI 使用您的 API 密钥登录
wandb login
指标记录器
您需要做的唯一更改是将指标记录器添加到您的配置中。Weights & Biases 将为您记录指标和模型检查点。
# enable logging to the built-in WandBLogger
metric_logger:
_component_: torchtune.training.metric_logging.WandBLogger
# the W&B project to log to
project: torchtune
我们会自动从您正在运行的食谱中获取配置并将其记录到 W&B。您可以在 W&B 概述选项卡中找到它,以及 Files
选项卡中的实际文件。
提示:如果您的作业崩溃或在未清理资源的情况下退出,您可能会在后台看到 wandb 进程残留。要杀死这些残留进程,可以使用类似 ps -aux | grep wandb | awk '{ print $2 }' | xargs kill
的命令。
注意
单击此示例 项目以查看 W&B 工作区。用于训练模型的配置可以在 此处 找到。
将模型检查点记录到 W&B
您还可以通过修改所需脚本 save_checkpoint
方法将模型检查点记录到 W&B。
建议的方法类似于以下内容
def save_checkpoint(self, epoch: int) -> None:
...
## Let's save the checkpoint to W&B
## depending on the Checkpointer Class the file will be named differently
## Here is an example for the full_finetune case
checkpoint_file = Path.joinpath(
self._checkpointer._output_dir, f"torchtune_model_{epoch}"
).with_suffix(".pt")
wandb_at = wandb.Artifact(
name=f"torchtune_model_{epoch}",
type="model",
# description of the model checkpoint
description="Model checkpoint",
# you can add whatever metadata you want as a dict
metadata={
training.SEED_KEY: self.seed,
training.EPOCHS_KEY: self.epochs_run,
training.TOTAL_EPOCHS_KEY: self.total_epochs,
training.MAX_STEPS_KEY: self.max_steps_per_epoch,
}
)
wandb_at.add_file(checkpoint_file)
wandb.log_artifact(wandb_at)