Pytorch使用TensorBoard测试

参考图书：深入浅出Pytorch
代码：

import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
from sklearn.datasets import load_boston


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):
        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)
write = SummaryWriter()

for step in range(10000):
    # 迭代
    predict = lm(data)
    loss = criterion(predict, target)
    write.add_scalar("Loss/train", loss, step)
    write.add_histogram("Param/weight", lm.weight, step)
    write.add_histogram("Param/bias", lm.bias, step)

    if step and step % 500 == 0:
        print("Loss:{:.3f}".format(loss.item()))

    optim.zero_grad()
    loss.backward()
    optim.step()

运行之后终端会出现runs子文件，之后输入：

tensorboard --logdir ./runs

最终会在本机6006端口会现实tensorboard可视化结果：