Mxnet 网络可视化（MobileNetV2）

红太狼 2022-04-10 01:48 308阅读 0赞

import matplotlib
    matplotlib.use('Agg')
    
    import argparse,time,logging
    import mxnet as mx
    import numpy as np
    from mxnet import gluon,nd
    from mxnet.gluon import nn
    from mxnet.gluon.data.vision import transforms
    
    from gluoncv.data import imagenet
    from gluoncv.utils import makedirs, TrainingHistory
    
    import os
    from mxnet.context import cpu
    from mxnet.gluon.block import HybridBlock
    from mxnet.gluon.contrib.nn import HybridConcurrent
    import multiprocessing
    import logging
    logging.basicConfig(level=logging.INFO)
    
    ##This block contains definition for Mobilenet v2
    
    # Helpers
    class RELU6(nn.HybridBlock):
        """Relu6 used in MobileNetV2."""
    
        def __init__(self, **kwargs):
            super(RELU6, self).__init__(**kwargs)
    
        def hybrid_forward(self, F, x):
            return F.clip(x, 0, 6, name="relu6")
    
    
    def _add_conv(out, channels=1, kernel=1, stride=1, pad=0,
                  num_group=1, active=True, relu6=False):
        out.add(nn.Conv2D(channels, kernel, stride, pad, groups=num_group, use_bias=False))
        out.add(nn.BatchNorm(scale=True))
        if active:
            out.add(RELU6() if relu6 else nn.Activation('relu'))
    
    
    def _add_conv_dw(out, dw_channels, channels, stride, relu6=False):
        _add_conv(out, channels=dw_channels, kernel=3, stride=stride,
                  pad=1, num_group=dw_channels, relu6=relu6)
        _add_conv(out, channels=channels, relu6=relu6)
        
    
    class LinearBottleneck(nn.HybridBlock):
        r"""LinearBottleneck used in MobileNetV2 model from the
        `"Inverted Residuals and Linear Bottlenecks:
          Mobile Networks for Classification, Detection and Segmentation"
        <https://arxiv.org/abs/1801.04381>`_ paper.
        Parameters
        ----------
        in_channels : int
            Number of input channels.
        channels : int
            Number of output channels.
        t : int
            Layer expansion ratio.
        stride : int
            stride
        """
    
        def __init__(self, in_channels, channels, t, stride, **kwargs):
            super(LinearBottleneck, self).__init__(**kwargs)
            self.use_shortcut = stride == 1 and in_channels == channels
            with self.name_scope():
                self.out = nn.HybridSequential()
    
                _add_conv(self.out, in_channels * t, relu6=True)
                _add_conv(self.out, in_channels * t, kernel=3, stride=stride,
                          pad=1, num_group=in_channels * t, relu6=True)
                _add_conv(self.out, channels, active=False, relu6=True)
    
        def hybrid_forward(self, F, x):
            out = self.out(x)
            if self.use_shortcut:
                out = F.elemwise_add(out, x)
            return out
        
    # Net
    class MobileNetV2(nn.HybridBlock):
        r"""MobileNetV2 model from the
        `"Inverted Residuals and Linear Bottlenecks:
          Mobile Networks for Classification, Detection and Segmentation"
        <https://arxiv.org/abs/1801.04381>`_ paper.
        Parameters
        ----------
        multiplier : float, default 1.0
            The width multiplier for controling the model size. The actual number of channels
            is equal to the original channel size multiplied by this multiplier.
        classes : int, default 1000
            Number of classes for the output layer.
        """
    
        def __init__(self, multiplier=1.0, classes=1000, **kwargs):
            super(MobileNetV2, self).__init__(**kwargs)
            with self.name_scope():
                self.features = nn.HybridSequential(prefix='features_')
                with self.features.name_scope():
                    _add_conv(self.features, int(32 * multiplier), kernel=3,
                              stride=2, pad=1, relu6=True)
    
                    in_channels_group = [int(x * multiplier) for x in [32] + [16] + [24] * 2
                                         + [32] * 3 + [64] * 4 + [96] * 3 + [160] * 3]
                    channels_group = [int(x * multiplier) for x in [16] + [24] * 2 + [32] * 3
                                      + [64] * 4 + [96] * 3 + [160] * 3 + [320]]
                    ts = [1] + [6] * 16
                    strides = [1, 2] * 2 + [1, 1, 2] + [1] * 6 + [2] + [1] * 3
    
                    for in_c, c, t, s in zip(in_channels_group, channels_group, ts, strides):
                        self.features.add(LinearBottleneck(in_channels=in_c, channels=c,
                                                           t=t, stride=s))
    
                    last_channels = int(1280 * multiplier) if multiplier > 1.0 else 1280
                    _add_conv(self.features, last_channels, relu6=True)
    
                    self.features.add(nn.GlobalAvgPool2D())
    
                self.output = nn.HybridSequential(prefix='output_')
                with self.output.name_scope():
                    self.output.add(
                        nn.Conv2D(classes, 1, use_bias=False, prefix='pred_'),
                        nn.Flatten()
                    )
    
        def hybrid_forward(self, F, x):
                x = self.features(x)
                x = self.output(x)
                return x
    
    # Constructor
    def get_mobilenet_v2(multiplier, **kwargs):
        r"""MobileNetV2 model from the
        `"Inverted Residuals and Linear Bottlenecks:
          Mobile Networks for Classification, Detection and Segmentation"
        <https://arxiv.org/abs/1801.04381>`_ paper.
        Parameters
        ----------
        multiplier : float
            The width multiplier for controling the model size. Only multipliers that are no
            less than 0.25 are supported. The actual number of channels is equal to the original
            channel size multiplied by this multiplier.
        """
        net = MobileNetV2(multiplier, **kwargs)
        return net
    
    def mobilenet_v2_1_0(**kwargs):
        r"""MobileNetV2 model from the
        `"Inverted Residuals and Linear Bottlenecks:
          Mobile Networks for Classification, Detection and Segmentation"
        <https://arxiv.org/abs/1801.04381>`_ paper.
        """
        return get_mobilenet_v2(1.0, **kwargs)
    
    def mobilenet_v2_0_5(**kwargs):
        r"""MobileNetV2 model from the
        `"Inverted Residuals and Linear Bottlenecks:
          Mobile Networks for Classification, Detection and Segmentation"
        <https://arxiv.org/abs/1801.04381>`_ paper.
        """
        return get_mobilenet_v2(0.5, **kwargs)
    models = {  
                'mobilenetv2_1.0': mobilenet_v2_1_0,
                'mobilenetv2_0.5': mobilenet_v2_0_5
             }
    
    kwargs={'classes':1000}
    # Retireve gluon model
    net = models['mobilenetv2_1.0'](**kwargs)
    
    ctx = mx.cpu(0)
    net.hybridize()
    net.initialize(mx.init.MSRAPrelu(), ctx=ctx)
    
    # 可视化网络
    x = mx.sym.var('data')
    sym = net(x)
    mx.viz.plot_network(sym,node_attrs={"shape":'oval',"fixedsize":'false'}).view() #pycharm使用必须加view()