TensorBoard的安装

要求Pytorch版本必须在1.2.0以上。

使用下面的命令安装：

1
from torch.utils.tensorboard import SummaryWriter

TensorBoard的代码调用

（1）导入包，并创建TensorBoard回调对象

1
2
from tensorflow.keras.callbacks import TensorBoard
writer = SummaryWriter("logs/learning_rate_scheduler") #指定TensorBoard日志目录

（2）在模型的训练过程中导入回调

1
2
3
4
5
6
7
8
9
10
11
global_step = 0 # 初始化 global_step 为 0

for epoch in range(num_epochs):
    for batch_idx, (data, target) in enumerate(train_loader):
        # 训练过程
        ...
        
        # 将学习率和训练损失添加到 TensorBoard
        writer.add_scalar('Train/Loss', loss, global_step=global_step)
        writer.add_scalar('Train/Learning_Rate', lr, global_step=global_step)
		global_step += 1  # 为每个batch更新 global_step 计数器

查看曲线

训练开始后，打开一个终端，输入下面的命令

1
tensorboard --logdir logs/learning_rate_scheduler

然后打开浏览器的http://localhost:6006/页面，就可以看到曲线。