Pytorch使用Tensorboard记录loss曲线 (Tensorboard学习二)

关于Tensorboard的基本使用方法可以参考：

对于一个基本模型：

import torch
import torch.nn as nn


class LinearRegressionModel(nn.Module):
    def __init__(self, input_shape, output_shape):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_shape, output_shape)

    def forward(self, x):
        out = self.linear(x)
        return out


def train_model():
    x_train = torch.randn(100, 4)  # 生成100个4维的随机数，作为训练集的 X
    y_train = torch.randn(100, 1)  # 作为训练集的label

    model = LinearRegressionModel(x_train.shape[1], 1)
    epochs = 100  # 迭代1000次
    learning_rate = 0.01  # 学习率
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)  # 优化函数
    criterion = nn.MSELoss()  # Loss使用MSE值，目标是使MSE最小

    for epoch in range(epochs):
        epoch += 1
        optimizer.zero_grad()
        outputs = model(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
    # 预测
    model.eval()
    predicted = model(torch.randn(100, 4)).data.numpy()
    print(predicted)


if __name__ == __main__:
    train_model()

改写代码，添加绘图代码，包括三步:

1. 导包：from torch.utils.tensorboard import SummaryWriter
2. 定义writer：writer = SummaryWriter()
3. 写入数据：writer.add_scalar("图像名称",y值, x值)
4. 关闭并写入：writer.close()

import torch
import torch.nn as nn


class LinearRegressionModel(nn.Module):
    def __init__(self, input_shape, output_shape):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_shape, output_shape)

    def forward(self, x):
        out = self.linear(x)
        return out


def train_model():
    x_train = torch.randn(100, 4)  # 生成100个4维的随机数，作为训练集的 X
    y_train = torch.randn(100, 1)  # 作为训练集的label

    model = LinearRegressionModel(x_train.shape[1], 1)
    epochs = 100  # 迭代1000次
    learning_rate = 0.01  # 学习率
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)  # 优化函数
    criterion = nn.MSELoss()  # Loss使用MSE值，目标是使MSE最小

    for epoch in range(epochs):
        epoch += 1
        optimizer.zero_grad()
        outputs = model(x_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        writer.add_scalar("loss", loss.detach(), epoch)  # 第三步，绘图
    # 预测
    model.eval()
    predicted = model(torch.randn(100, 4)).data.numpy()
    print(predicted)
    writer.close()  # 第4步，写入关闭


if __name__ == __main__:
    from torch.utils.tensorboard import SummaryWriter # 第1步

    writer = SummaryWriter(log_dir="summary_pic")  # 第二步，确定保存的路径，会保存一个文件夹，而非文件
    # tensorboard --logdir=summary_pic
    train_model()

使用：

tensorboard --logdir=summary_pic

在本地的http://localhost:6006/打开可以得到相关图像：