关于Tensorboard的基本使用方法可以参考:Pytorch 配置 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()
改写代码,添加绘图代码,包括三步:
- 导包:
from torch.utils.tensorboard import SummaryWriter
- 定义writer:
writer = SummaryWriter()
- 写入数据:
writer.add_scalar("图像名称",y值, x值)
- 关闭并写入:
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/
打开可以得到相关图像: