模型结构可视化分析

模型结构可视化分析

很多时候需要得知自己所构建模型的结构，参数，所占存储容量等，下面介绍几个可视化的参数工具

1、pytorch-summary（包含每一层的输入输出形状，参数量，以及所占存储量大小）

（1）pip install torchsummary 安装相关库

（2）举例说明：

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class simpleconv3(nn.Module):
    def __init__(self,nclass):
        super(simpleconv3,self).__init__()
        self.conv1 = nn.Conv2d(3, 12, 3, 2)
        self.bn1 = nn.BatchNorm2d(12)
        self.conv2 = nn.Conv2d(12, 24, 3, 2)
        self.bn2 = nn.BatchNorm2d(24)
        self.conv3 = nn.Conv2d(24, 48, 3, 2)
        self.bn3 = nn.BatchNorm2d(48)
        self.fc1 = nn.Linear(48 * 5 * 5 , 1200)
        self.fc2 = nn.Linear(1200 , 128)
        self.fc3 = nn.Linear(128 , nclass)

    def forward(self , x):
        # print(x.shape)
        x = F.relu(self.bn1(self.conv1(x)))
        # print(x.shape)
        #print "bn1 shape",x.shape
        x = F.relu(self.bn2(self.conv2(x)))
        # print(x.shape)
        x = F.relu(self.bn3(self.conv3(x)))
        # print(x.shape)
        x = x.view(-1 , 48 * 5 * 5) 
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

这是所构建的一个简单的三层网络

#参数量分析工具pytorch-summary
from torchsummary import summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#有GPU则使用GPU没有则使用cpu
model = simpleconv3(2).to(device)#这里是二分类任务
summary(model,(3,48,48))  #输入模型以及图片的尺寸

这里是显示的结果：

能够看到每层的类型

Layer：层的类型
Output Shape: 形状
param：每层的参数量
Total param：总参数量
Trainable params：可训练的参数量
Forward/backward pass size：一次需要的内存大小，可用来估计最佳bach_size

2、Netron（可视化网络结构，权重尺寸与大小）

这里其实是有两种方式的，一种是直接用网页打开，另一种需要下载软件到本地

（1）网页版

点击这里

再点击第一个按键，选择你要可视化的模型，这里可能会出现下面这种情况。

将自定义文件换为所有文件（如下图），点击打开，这里的网络结构与上面的是一样的。

（这里解释一下，上图中的偏差和权重都是4维的是因为之前做分类任务时做的是四分类，所以模型是四分类的模型，但结构是类似的。）

可以看见该网络结构的权重和偏置，也可以看见每层的详细参数，是一个很方便的工具。

（2）下载软件进行可视化。（这里就不多做介绍了，效果上基本是一致的，只是显示方式是直立的图）

（3）onnx格式，在保存模型时可以将模型保存为onnx格式的模型，如下示例代码：

#模型结构可视化工具Netron
#pytorch ONNX格式导出

mynet.load_state_dict(torch.load('./models/model.ckpt',map_location=lambda storage,loc:storage))
mynet.train(False)
dummy_input = torch.randn((1,3,48,48))
torch.onnx.export(mynet,dummy_input,'model.onnx',verbose=False)

下面这些是对上面代码行的注解

'''
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=torch.device('cpu'))
# Load all tensors onto the CPU, using a function
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
'''

'''
torch.onnx.export(model, args,export_params=True, verbose=False)
model (torch.nn.Module) – 要导出的模型.
args (tuple of arguments) – 模型的输入, 任何非Tensor参数都将硬编码到导出的模型中；任何Tensor参数都将成为导出的模型的输入，并按照他们在args中出现的顺序输入。因为export运行模型，所以我们需要提供一个输入张量x。只要是正确的类型和大小，其中的值就可以是随机的。请注意，除非指定为动态轴，否则输入尺寸将在导出的ONNX图形中固定为所有输入尺寸。在此示例中，我们使用输入batch_size 1导出模型，但随后dynamic_axes 在torch.onnx.export()。因此，导出的模型将接受大小为[batch_size，3、100、100]的输入，其中batch_size可以是可变的。
export_params (bool, default True) – 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False.
verbose (bool, default False) - 如果指定，我们将打印出一个导出轨迹的调试描述。
'''

显示结果：

与上图具有明显的不同，可以根据不同的使用需求进行选择。

3、可视化工具Graphviz（对网络结构进行可视化，可视化的结构以PDF格式保存）

这里的代码可以说是一套标准的格式了，经过简单的修改便可用在其他网络结构中。

if __name__ == '__main__':
    import torch
    from torch.autograd import Variable
    from visualize import  make_dot
    x = Variable(torch.randn(1,3,48,48))
    model = simpleconv3(4)
    y = model(x)
    g = make_dot(y)
    g.view()

Variable.py文件：（不需要修改）

from graphviz import Digraph
import torch
from torch.autograd import Variable

def make_dot(var, params=None):
    """ Produces Graphviz representation of PyTorch autograd graph
    Blue nodes are the Variables that require grad, orange are Tensors
    saved for backward in torch.autograd.Function
    Args:
        var: output Variable
        params: dict of (name, Variable) to add names to node that
            require grad (TODO: make optional)
    """
    if params is not None:
        assert isinstance(params.values()[0], Variable)
        param_map = {id(v): k for k, v in params.items()}

    node_attr = dict(style='filled',
                     shape='box',
                     align='left',
                     fontsize='12',
                     ranksep='0.1',
                     height='0.2')
    dot = Digraph(node_attr=node_attr, graph_attr=dict(size="12,12"))
    seen = set()

    def size_to_str(size):
        return '('+(', ').join(['%d' % v for v in size])+')'

    def add_nodes(var):
        if var not in seen:
            if torch.is_tensor(var):   #判断是否是节点
                dot.node(str(id(var)), size_to_str(var.size()), fillcolor='orange')
            elif hasattr(var, 'variable'):    #边
                u = var.variable
                name = param_map[id(u)] if params is not None else ''
                node_name = '%s\n %s' % (name, size_to_str(u.size()))
                dot.node(str(id(var)), node_name, fillcolor='lightblue')
            else:
                dot.node(str(id(var)), str(type(var).__name__))
            seen.add(var)
            if hasattr(var, 'next_functions'):
                for u in var.next_functions:
                    if u[0] is not None:
                        dot.edge(str(id(u[0])), str(id(var)))
                        add_nodes(u[0])
            if hasattr(var, 'saved_tensors'):
                for t in var.saved_tensors:
                    dot.edge(str(id(t)), str(id(var)))
                    add_nodes(t)
    add_nodes(var.grad_fn)
    return dot

最终结果图（一个pdf文件）：

关于可视化工具Graphviz可能会出现报错，这里需要在win10上安装相应的Graphviz，具体的可以直接百度，网上教程很多的。以上三种查看网络结构的方法各有各的好处，可以根据自己的情况进行选择。