pytorch中可视化网络结构方法(1)

1.安装两个包

pip install tensorboardX

pip install torchsummary

2.实例.

import torch
from torch import nn
from torch.nn import functional as F
from tensorboardX import SummaryWriter

class MyNet(nn.Module):  # 为了方便自定义的三层网络
      def __init__(self):
          super(MyNet, self).__init__()
          self.linear1 = nn.Linear(10,16)
          self.linear2 = nn.Linear(16,10)
          self.linear3 = nn.Linear(10,5)

      def forward(self, x):
          x = F.relu(self.linear1(x),inplace=True)
          x = F.relu(self.linear2(x),inplace=True)
          x = F.relu(self.linear3(x), inplace=True)
          return x

x = torch.randn(1,10)
net = MyNet()

with SummaryWriter(comment='MyNet') as w:
    w.add_graph(net, x)  

3.程序运行完成后，会生成一个run文件件，找到并打开该文件夹，在该文件夹中运行

tensorboard --logdir Jan08_17-43-58_DESKTOP-K6BQG5GMyNet

# 注意要与自己的文件夹名字相同

Windows在指定文件夹中打开DOS命令窗口的方法：

在文件夹中按住shift键右击：在此处打开PowerShell窗口，在窗口中输入start cmd即可，最后将DOS窗口中的地址复制到浏览器即可