我们在搭建网络模型的时候,通常希望可以对自己搭建好的网络模型有一个比较好的直观感受,从而更好地了解网络模型的结构,Tensorboard工具的使用就给我们提供了方便的途径
Tensorboard是由Google公司开源的一款可视化工具,是TensorFlow的一个附属组件,但在pytorch项目中也可以使用。
它有以下主要功能:
这篇博客主要介绍Tensorboard可视化网络模型的功能
我们搭建一个简单的神经网络,依赖的库环境
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
搭建网络模型
class Net(nn.Module):
def __init__(self,input_dim,layer1_dim,layer2_dim,output_dim):
super(Net,self).__init__()
self.flatten = nn.Flatten()
self.layer1 = nn.Sequential(nn.Linear(input_dim,layer1_dim),nn.ReLU())
self.layer2 = nn.Sequential(nn.Linear(layer1_dim,layer2_dim),nn.ReLU())
self.out = nn.Sequential(nn.Linear(layer2_dim,output_dim),nn.Softmax(dim=-1))
def forward(self,x):
x = self.flatten(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.out(x)
return x
# 初始化网络中的值
input_dim,layer1_dim,layer2_dim,output_dim=32*32,512,128,10
model = Net(input_dim,layer1_dim,layer2_dim,output_dim)
网络由三层全连接层组成,输入的数据形状为
# 定义输入模型的数据,1表示批次,这里可以忽略
input_data = torch.rand(1,32,32)
# 模型输出的数据
output_data = model(input_data)
接下俩就是使用SummaryWriter创建日志保存搭建好的网络模型
with SummaryWriter(log_dir=r"D:\CSDN_point\12_22\logs", comment="Net") as w:
w.add_graph(model, input_data)
log_dir参数就是日志文件的本地保存路径,comment就是日志的备注,add_graph()传入网络模型和输入数据,运行后就会在指定路径上生成对应的文件,打开log文件所在的文件位置,在顶部路径上输入cmd,打开命令行窗口
在命令行窗口输入
tensorboard --logdir logs
复制返回的网址在浏览器打开,就可以得到对应的网络可视化结果了
欢迎大家讨论交流~