使用TensorBoard可视化训练过程

TensorBoard是一个用于追踪、可视化、分析模型训练过程和训练结果的工具，它提供了多种可视化功能，可以与PyTorch、TensorFlow、Keras、Huggingface transformers、ModelScope等机器学习框架一起使用，帮助用户了解模型的训练过程和性能。

PAI提供了TensorBoard服务，支持用户在PAI创建TensorBoard应用，用于查看训练作业输出的TensorBoard日志。

本文档将以不同的机器学习框架为示例，展示如何在PAI使用TensorBoard追踪和可视化模型训练过程。

安装和配置SDK

我们需要首先安装PAI Python SDK以运行本示例。

!python -m pip install --upgrade alipai

SDK需要配置访问阿里云服务需要的AccessKey，以及当前使用的工作空间和OSS Bucket。在PAI SDK安装之后，通过在 命令行终端 中执行以下命令，按照引导配置密钥、工作空间等信息。

# 以下命令，请在 命令行终端 中执行.

python -m pai.toolkit.config

我们可以通过以下代码验证配置是否已生效。

import pai
from pai.session import get_default_session

print(pai.__version__)

sess = get_default_session()

# 获取配置的工作空间信息
assert sess.workspace_name is not None
print(sess.workspace_name)

提交训练任务

我们首先需要准备训练脚本，使用将PyTorch的TensorBoard utility记录TensorBoard日志。

PyTorch提供的TensorBoard utilities的使用可以见文档： torch.utils.tensorboard 文档

!mkdir -p src

镜像里需要先安装TensorBoard，可以在训练目录中准备 requirements.txt 指定需要按照的第三方库。

%%writefile src/requirements.txt

tensorboard

%%writefile src/run.py

import os

import torch
from torch.utils.tensorboard import SummaryWriter


# 通过环境变量获取TensorBoard输出路径，默认为 /ml/output/tensorboard/
tb_log_dir = os.environ.get("PAI_OUTPUT_TENSORBOARD")
print(f"TensorBoard log dir: {tb_log_dir}")
writer = SummaryWriter(log_dir=tb_log_dir)

def train_model(iter):


    x = torch.arange(-5, 5, 0.1).view(-1, 1)
    y = -5 * x + 0.1 * torch.randn(x.size())

    model = torch.nn.Linear(1, 1)
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr = 0.1)

    for epoch in range(iter):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

if __name__ == "__main__":
    train_model(100)
    writer.flush()

from pai.estimator import Estimator
from pai.image import retrieve


est = Estimator(
    command="python run.py",
    source_dir="./src",
    image_uri=retrieve("PyTorch", "latest").image_uri,
    instance_type="ecs.c6.large",
)

est.fit(wait=False)

使用TensorBoard应用监控训练

在PAI启动一个TensorBoard应用，查看使用Estimator的训练作业写出的TensorBoard日志。

tb = est.tensorboard()

print(tb.app_uri)

使用完成之后，删除TensorBoard应用

tb.delete()