pytorch网络结构可视化（torchsummary+tensorboardX+网络多输通道入情况+多分支网络）

1.单通道输入网络

单通道输入的情况大致有以下两种结构：

结构1

只有一条路可以走

结构2

输入为一条路，输出为多条路

以上两种的输入只有一个input，这种是经常遇到的情况。

实例01

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        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)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)
if __name__ == '__main__':
	Net =Net()
	'''torchsummary 打印网络结构'''
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
	model = Net.to(device)
	summary(model, (1, 28, 28))#（1，28，28）对应（c,h,w）,不用指定batch参数
	'''tensorboardX生成日志文件'''
	dummy_input01 = torch.rand(10, 1, 28, 28)  # 假设输入10张1*28*28的图片
	with SummaryWriter(comment='Net') as w:
	    w.add_graph(Net, (dummy_input01))
	# tensorboard.exe --logdir=F:\code\CNN02\runs\Jul20_09-29-44_LAPTOP-DTN3ORVQCNN02 --host=127.0.0.1

torchsummary 打印的网络结构：
特别注意：

summary(model, (1, 28, 28))

在summary指令中设置shape为（1,28,28）即大小为28*28通道数为1的图片，让其模拟作为网络的输入，使得网络流通，以展示网络结构。

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 10, 24, 24]             260
            Conv2d-2             [-1, 20, 8, 8]           5,020
         Dropout2d-3             [-1, 20, 8, 8]               0
            Linear-4                   [-1, 50]          16,050
            Linear-5                   [-1, 10]             510
================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.06
Params size (MB): 0.08
Estimated Total Size (MB): 0.15
----------------------------------------------------------------

tensorboardX展示网络结构：
特别注意

w.add_graph(Net, (dummy_input01))

dummy_input01 = torch.rand(10, 1, 28, 28) ，同样的，设置shape为（1,28,28）即大小为28*28通道数为1的图片，让其模拟作为网络的输入，使得网络流通，以展示网络结构。不一样的，多了10，即网络的真实输入shape(N,C,H,W)。

在命令框里输入指令（logdir=之后，–host=127.0.0.1之前，是生成的包含网络结构的日志文件地址）：

tensorboard.exe --logdir=F:\code\CNN02\runs\Jul20_09-29-44_LAPTOP-DTN3ORVQCNN02 --host=127.0.0.1

回车后，命令框打印出：

TensorBoard 2.0.2 at http://127.0.0.1:6006/ (Press CTRL+C to quit)

将*http://127.0.0.1:6006/*在浏览器中打开，即可查看网络结构：

2.多通道输入网络

结构

输入是多条路，这就意味着两条路里都要给定输入，即多通道输入。

实例02（只用了卷积层进行演示）

import torch
import torch.nn as nn
from torchsummary import summary
from tensorboardX import SummaryWriter

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 3, kernel_size=3,stride=1,padding=1,)
        self.conv2 = nn.Conv2d(1, 3, kernel_size=3,stride=1,padding=1,)
        self.conv3 = nn.Conv2d(6, 12, kernel_size=3, stride=1, padding=1)
        self.conv4 = nn.Conv2d(12,1, kernel_size=3, stride=1, padding=1)
    def forward(self, x1,x2):
        x1 = self.conv1(x1)
        x2 = self.conv2(x2)
        layer_merged = torch.cat((x1,x2),1)
        out = self.conv3(layer_merged)
        out = self.conv4(out)
        return out
if __name__ == '__main__':
    Net =Net()
    '''torchsummary 打印网络结构'''
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = Net.to(device)
    summary(model, [(1, 28, 28),(1, 28, 28)])
    '''tensorboardX生成日志文件'''
    dummy_input01 = torch.rand(10, 1, 28, 28)  # 假设输入10张1*28*28的图片
    dummy_input02 = torch.rand(10, 1, 28, 28)  # 假设输入10张1*28*28的图片
    with SummaryWriter(comment='Net') as w:
        w.add_graph(Net, (dummy_input01,dummy_input02))
    # tensorboard.exe --logdir=F:\code\CNN02\runs\Jul20_09-29-44_LAPTOP-DTN3ORVQCNN02 --host=127.0.0.1

torchsummary 打印的网络结构：
特别注意：

 summary(model, [(1, 28, 28),(1, 28, 28)])

因为输入有两个通道即需要两路的input，注意summary参数的给定，summary(model, [(1, 28, 28),(1, 28, 28)])。

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 3, 28, 28]              30
            Conv2d-2            [-1, 3, 28, 28]              30
            Conv2d-3           [-1, 12, 28, 28]             660
            Conv2d-4            [-1, 1, 28, 28]             109
================================================================
Total params: 829
Trainable params: 829
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 2.34
Forward/backward pass size (MB): 0.11
Params size (MB): 0.00
Estimated Total Size (MB): 2.46
----------------------------------------------------------------

tensorboardX展示网络结构：
特别注意

w.add_graph(Net, (dummy_input01,dummy_input02))

dummy_input01 = torch.rand(10, 1, 28, 28)  # 假设输入10张1*28*28的图片
dummy_input02 = torch.rand(10, 1, 28, 28)  # 假设输入10张1*28*28的图片

同样的，注意参数的设置。
生成日志文件等操作不再赘述，在浏览器中查看网络结构如下：

3.参考链接：

1.summary示例
2.tensorboardx显示网络架构官方文档
3.详解PyTorch项目使用TensorboardX进行训练可视化
4.tensorboard生成的网址打不开的解决方法