这里以add_graph为例
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from tensorboardX import SummaryWriter
class Net1(nn.Module):
def __init__(self):
super(Net1, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)#padding=0
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
self.bn = nn.BatchNorm2d(20)
def forward(self, x):
x = F.max_pool2d(self.conv1(x), 2)#input,kernel_size,padding=0,stride=2
x = F.relu(x) + F.relu(-x)
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = self.bn(x)
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
x = F.softmax(x, dim=1)
return x
dummy_input = Variable(torch.rand(13, 1, 28, 28))#batch_size,in_channels,height,width
model = Net1()
with SummaryWriter(comment='Net1') as w:
w.add_graph(model, (dummy_input, ),True)
执行程序之后,在当前工程中生成了一个runs文件夹,文件夹中存放了一个保存的模型,这里是不可读的,需要再执行一条命令。
这里前提是先安装一个库pip install tensorboard
定位到runs文件夹的上一级,打开终端输入命令
tensorboard --logdir runs
幸运的话,运行得到的链接直接复制到浏览器就可以看到生成的图了。
如果提示电脑拒绝连接,则换成下面这条命令:
tensorboard --logdir runs --host=127.0.0.1
我是通过这条命令才能看到下面的图
如果觉得每次输入的命令太长,可以这样做:
终端打开bashrc~/.bashrc
添加一条语句tb=tensorboard --logdir
参考文章:
https://blog.youkuaiyun.com/weixin_44135282/article/details/86156961