PyTorch 实现可视化需要借助 TensorBoard 包。
参考：https://pytorch.org/docs/stable/tensorboard.html

安装

pip install tensorboard

使用

通过如下代码进行记录：

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter() # 实例化一个 SummaryWriter，无参数输入则默认保存在py文件同级目录下；也可以在括号中指定保存路径
writer.add_scalars(main_tag, tag_scalar_dict, global_step=None, walltime=None) # 标量可视化
writer.add_histogram(tag, values, global_step=None, bins='tensorflow', walltime=None, max_bins=None) # 直方图可视化
writer.add_image(tag, img_tensor, global_step=None, walltime=None, dataformats='CHW') # 图片可视化
writer.add_figure(tag, figure, global_step=None, close=True, walltime=None) # matplotlib 绘图可视化
writer.add_video(tag, vid_tensor, global_step=None, fps=4, walltime=None) # 视频可视化，需要有 moviepy 包
writer.add_audio(tag, snd_tensor, global_step=None, sample_rate=44100, walltime=None) # 音频可视化
writer.add_text(tag, text_string, global_step=None, walltime=None) # 文本可视化
writer.add_graph(model, input_to_model=None, verbose=False) # 计算图可视化
writer.add_hparams(hparam_dict=None, metric_dict=None) # 超参数可视化
...
writer.close()

通过在系统命令控制窗口输入：
tensorboard --logdir=E:\jupyter-notebook\2020_DeepLearning\code_trainning\runs
可打开 tensorboard 面板进行查看。

tensorboard 面板

注：E:\jupyter-notebook\2020_DeepLearning\code_trainning\runs为我的py文件的同级目录文件夹。

示例1：显示计算图

import torch
import torch.nn.functional as F 
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
# 生成输入数据
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1) # x data(tensor),shape=(100,1)
#搭建神经网络
class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_output):
        super(Net, self).__init__() 
        self.hidden = torch.nn.Linear(n_feature,n_hidden)
        self.predict = torch.nn.Linear(n_hidden,n_output)
        
    def forward(self,x):
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        return x
net = Net(n_feature=1,n_hidden=10,n_output=1)
# 写入 SummaryWriter
writer.add_graph(net, x) # 计算图可视化
writer.close()

通过在系统命令控制窗口输入：
tensorboard --logdir=E:\jupyter-notebook\2020_DeepLearning\code_trainning\runs
打开 tensorboard 面板。

计算图