在PyTorch中使用TensorBoard

在机器学习中，要改进模型的某些参数，我们通常需要对其进行衡量。TensorBoard 是用于提供机器学习工作流期间所需测量和呈现的工具。它使我们能够跟踪实验指标（例如损失和准确率），呈现模型计算图，将嵌入向量投影到较低维度的空间等。

1.安装

pip install tensorboard

2.TensorBoard使用

2.1创建SummaryWriter实例

import torch
from torch.utils.tensorboard import SummaryWriter

# 记录的内容保存在“./log”里面，log里面会生成events.out.tfevents的文件
writer = SummaryWriter('./log')  

2.2利用add_scalar()记录metrics

了解损失等关键指标以及它们在训练过程中的变化是非常重要的。标量有助于保存每个训练步骤的损失值，或者每个epoch之后的精度。

使用 add_scalar(tag, scalar_value, global_step=None, walltime=None)记录标量值。

# 用法示例：
writer.add_scalar('loss', loss, epoch)
writer.add_scalar('accuracy', accuracy, epoch)

'''
* @name: tensorboard_test.py
* @description:创建一个线性回归模型，并使用add_scalar记录损失值
'''
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)

def train_model(total_epoch):
    for epoch in range(total_epoch):
        y1 = model(x)
        loss = criterion(y1, y)
        writer.add_scalar("Loss/train", loss, epoch)  # <-- 注意看这里！！！
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
train_model(10)

writer.flush()  # 调用flush()方法确保所有挂起的事件都已写入磁盘

2.3关闭Writer

最后调用 close() 方法来关闭Writer，释放资源。

writer.close()

2.4启动TensorBoard

最后，启动TensorBoard，指定上面我们使用的根日志目录。当不需要可视化时，在终端按下Ctrl+C关闭TensorBoard即可。

# 第一种写法
tensorboard --logdir=./log --host=127.0.0.1 --port=6006 --bind_all
# 第二种写法
tensorboard --logdir ./log --host 127.0.0.1 --port 6006 --bind_all

# 常用启动命令（只需给定--logdir参数，其余默认）
tensorboard --logdir ./log

• 参数--logdir指向TensorBoard将在其中查找它可以显示的事件文件的目录。TensorBoard将递归遍历以logdir为根的目录结构，寻找.*tfevents.*文件。
• 参数--host表示IP地址，默认为127.0.0.1
• 参数--port表示端口号，默认为6006
• 参数--bind_all指定 TensorBoard 绑定所有网络接口，而不仅仅是本地接口。默认情况下，TensorBoard 仅在本地主机上运行，只能通过 localhost 访问。如果添加 --bind_all 参数，就可以从网络中的其他设备访问 TensorBoard。这在需要远程访问时非常有用，例如在服务器上训练模型并从个人电脑上查看训练过程。

该仪表板显示了损失随着epoch的变化，还可以使用它来跟踪训练速度、学习率和其他标量值。

3.本地连接服务器使用TensorBoard

由于我们通常是在服务器上运行项目，那么TensorBoard生成的events.out.tfevents格式文件是保存在服务器上面的，如果想要观测 metrics 就需要将文件下载到本地，然后启动TensorBoard才可以在本地查看。为了解决这一问题，我们可以使用SSH隧道（也称SSH端口转发），将远程服务器的TensorBoard转发到本地。具体实现如下：

3.1方法一：使用SSH命令进行本地端口转发

ssh -L [本地端口]:[远程地址]:[远程端口] [用户名]@[远程服务器IP]

# example:
ssh -L 6006:127.0.0.1:6006 root@192.168.1.224

解释：这条命令将本地的6006端口转发到远程服务器192.168.1.224的6006端口上。这样我们就可以通过在本地访问127.0.0.1:6006来访问远程服务器的127.0.0.1:6006。

3.2方法二：启动TensorBoard时添加参数–bind_all（推荐✅）

tensorboard --logdir ./log --bind_all

在本地通过(远程服务器IP地址):6006/来查看训练过程。

3.3方法三：在终端工具里面配置SSH隧道或端口转发

流行的终端工具肯定是Xshell，Xshell的隧道配置可以查看这篇博客【使用Xshell/ssh 连接 远程服务器使用Tensorboard】。但是我个人喜欢使用Tabby，Tabby是一款免费的颜值很高的终端工具，也推荐大家使用。

Tabby的端口转发的添加，在设置->配置和连接->端口里面进行配置添加。

参考

• How to use TensorBoard with PyTorch
• Visualizing Models, Data, and Training with TensorBoard
• torch.utils.tensorboard官方文档
• 使用TensorBoard可视化训练过程
• 在Pytorch中使用Tensorboard可视化训练过程

😃😃😃