1.目的

我们在训练模型的时候，不知道自己的超参数epochs设置的是够正确,能否训练出一个足够好的模型？
通过参看loss的走向，如果loss稳定在较低值附近，就可以判断该模型已经收敛。
可以通过plt.show自己绘制loss的走向图，当然也可以通过TensorBoard绘制loss的走向图，更加高大上。
TensorBoard绘图简单原理如下：1.将要绘制的数据保存为一个文件 2.用tensorboard命令读取文件并绘制图形

2.步骤

安装

1	pip install tensorboard

训练模型代码中保存loss值

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("logs")  # 写入文件存放的位置
for i in range(100):
    writer.add_scalar("y=2*x", 2 * i, i)  #  "y=2*x"为图片描述，2 * i 为纵坐标， i为step（横坐标）
writer.close()

启动TensorBoard

1	tensorboard --logdir=logs

–logdir ：读取文件所在的文件夹

3.举例

import torch
from torch.utils.tensorboard import SummaryWriter
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[2.0], [4.0], [6.0]])
# 终端中输入 tensorboard --logdir=./logs 启动tensorboard读取logs下文件，进行绘制
# epochs = 100  # 设置为100时模型还没有收敛
epochs = 1000


class LinearModel(torch.nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)

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


model = LinearModel()
loss = torch.nn.MSELoss(reduction='mean')
opmizier = torch.optim.SGD(model.parameters(), lr=0.001)
writer = SummaryWriter("logs")
for epoch in range(epochs):
    y_pred = model(x_data)
    l = loss(y_pred, y_data)
    print(l.item())
    writer.add_scalar('loss', l.item(), epoch)  # 使用tensorboard
    l.backward()
    opmizier.step()
    opmizier.zero_grad()

writer.close()
print("w:", model.linear.weight.item())
print("b:", model.linear.bias.item())

x_test = torch.tensor([4.0])
y_test = model(x_test)
print("y_pred:", y_test.item())