使用tensorboardX可视化神经网络

#导入需要模块
from tensorboardX import SummaryWriter
import torch
import torch.nn as nn
import torchvision

#创建一个模型
class StemBlock(nn.Module):
    def __init__(self):
        super(StemBlock,self).__init__()
        self.stem_1 = nn.Conv2d(3,32,3,2,1)
        self.stem_2a = nn.Conv2d(32,16,1,1)
        self.stem_2b = nn.MaxPool2d(2,2)
        self.stem_2c = nn.Conv2d(16,32,3,2,1)
        self.stem_3 = nn.Conv2d(64,32,1)


    def forward(self,x):
        stem_1_out = self.stem_1(x)
        stem_2a_out = self.stem_2a(stem_1_out)
        stem_2c_out = self.stem_2c(stem_2a_out)
        stem_2b_out = self.stem_2b(stem_1_out)
        print(stem_2b_out.shape)
        print(stem_2c_out.shape)
        out = self.stem_3(torch.cat((stem_2b_out,stem_2c_out),1))
        
        return out

#把模型保存为graph

input = torch.randn(1, 3, 224, 224)
model = StemBlock()
with SummaryWriter(log_dir='./logs/',comment='test') as writer:
    writer.add_graph(model,input)

运行成功会在文件里出现下图的log

在终端输入以下内容

tensorboard --logdir ./logs/

会出现一个链接，点开链接