pytorch的数据可视化需要先安装TensorBoard。如果使用Anaconda,则需要进入相应的虚拟环境安装该包。
pip install tensorflow-tensorboard==1.5.1
pip install tensorboard==1.14.0
注意,tensorboard的版本需要大于等于1.14.0,否则在pytorch环境下使用会报错。
下面以简单线性回归模型为例,使用线性回归模型进行经典的波士顿房价预测,并对训练损失过程进行可视化展示。
import torch
import torch.nn as nn
from sklearn.datasets import load_boston
from torch.utils.tensorboard import SummaryWriter
class LinearModel(nn.Module):
def __init__(self, ndim):
super(LinearModel, self).__init__()
self.ndim = ndim
self.weight = nn.Parameter(torch.randn(ndim, 1)) # 定义权重
self.bias = nn.Parameter(torch.randn(1)) # 定义偏置
def forward(self, x):
# y = Wx + b
return x.mm(self.weight) + self.bias
boston = load_boston()
lm = LinearModel(13)
criterion = nn.MSELoss()
optim = torch.optim.SGD(lm.parameters(), lr=1e-6)
data = torch.tensor(boston['data'], requires_grad=True, dtype=torch.float32)
target = torch.tensor(boston['target'], dtype=torch.float32)
writer = SummaryWriter() # 定义TensorBoard输出类
for step in range(10000):
predict = lm(data)
loss = criterion(predict, target)
writer.add_scalar('Loss/train', loss, step) # 输出损失函数
writer.add_histogram('Param/weight', lm.weight, step) # 输出权重直方图
if step and step % 1000 == 0:
print('Loss: {:.3f}'.format(loss.item()))
optim.zero_grad()
loss.backward()
optim.step()
训练完成后,会在项目所在目录生成一个runs文件夹,此文件夹下还有一个文件夹,具体的文件夹名字与训练开始时间、用户主机名称有关。
接下来打开终端cmd,cd到项目所在目录,例如笔者的项目目录为:
继续执行“tensorboard --logdir runs”,生成pytorch模型的可视化数据。
最后在浏览器中访问http://127.0.0.1:6006,即可查看pytorch模型的已定义训练结果可视化的内容。