TensorRT-Demo/yolo/yolo_to_onnx.py

# yolo_to_onnx.py
#
# Copyright 1993-2019 NVIDIA Corporation.  All rights reserved.
#
# NOTICE TO LICENSEE:
#
# This source code and/or documentation ("Licensed Deliverables") are
# subject to NVIDIA intellectual property rights under U.S. and
# international Copyright laws.
#
# These Licensed Deliverables contained herein is PROPRIETARY and
# CONFIDENTIAL to NVIDIA and is being provided under the terms and
# conditions of a form of NVIDIA software license agreement by and
# between NVIDIA and Licensee ("License Agreement") or electronically
# accepted by Licensee.  Notwithstanding any terms or conditions to
# the contrary in the License Agreement, reproduction or disclosure
# of the Licensed Deliverables to any third party without the express
# written consent of NVIDIA is prohibited.
#
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
# SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE.  IT IS
# PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
# NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
# DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
# NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
# LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
# SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
# DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
# WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
# ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
# OF THESE LICENSED DELIVERABLES.
#
# U.S. Government End Users.  These Licensed Deliverables are a
# "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
# 1995), consisting of "commercial computer software" and "commercial
# computer software documentation" as such terms are used in 48
# C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
# only as a commercial end item.  Consistent with 48 C.F.R.12.212 and
# 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
# U.S. Government End Users acquire the Licensed Deliverables with
# only those rights set forth herein.
#
# Any use of the Licensed Deliverables in individual and commercial
# software must include, in the user documentation and internal
# comments to the code, the above Disclaimer and U.S. Government End
# Users Notice.
#


import os
import sys
import argparse
from collections import OrderedDict

import numpy as np
import onnx
from onnx import helper, TensorProto


MAX_BATCH_SIZE = 1


def parse_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-c', '--category_num', type=int,
        help='number of object categories (obsolete)')
    parser.add_argument(
        '-m', '--model', type=str, required=True,
        help=('[yolov3-tiny|yolov3|yolov3-spp|yolov4-tiny|yolov4|'
              'yolov4-csp|yolov4x-mish|yolov4-p5]-[{dimension}], where '
              '{dimension} could be either a single number (e.g. '
              '288, 416, 608) or 2 numbers, WxH (e.g. 416x256)'))
    args = parser.parse_args()
    return args


def rreplace(s, old, new, occurrence=1):
    """Replace old pattern in the string with new from the right."""
    return new.join(s.rsplit(old, occurrence))


def is_pan_arch(cfg_file_path):
    """Determine whether the yolo model is with PAN architecture."""
    with open(cfg_file_path, 'r') as f:
        cfg_lines = [l.strip() for l in f.readlines()]
    yolos_or_upsamples = [l for l in cfg_lines
                            if l in ['[yolo]', '[upsample]']]
    yolo_count = len([l for l in yolos_or_upsamples if l == '[yolo]'])
    upsample_count = len(yolos_or_upsamples) - yolo_count
    assert yolo_count in (2, 3, 4)  # at most 4 yolo layers
    assert upsample_count == yolo_count - 1 or upsample_count == 0
    # the model is with PAN if an upsample layer appears before the 1st yolo
    return yolos_or_upsamples[0] == '[upsample]'


def get_output_convs(layer_configs):
    """Find output conv layer names from layer configs.

    The output conv layers are those conv layers immediately proceeding
    the yolo layers.

    # Arguments
        layer_configs: output of the DarkNetParser, i.e. a OrderedDict of
                       the yolo layers.
    """
    output_convs = []
    previous_layer = None
    for current_layer in layer_configs.keys():
        if previous_layer is not None and current_layer.endswith('yolo'):
            assert previous_layer.endswith('convolutional')
            activation = layer_configs[previous_layer]['activation']
            if activation == 'linear':
                output_convs.append(previous_layer)
            elif activation == 'logistic':
                output_convs.append(previous_layer + '_lgx')
            else:
                raise TypeError('unexpected activation: %s' % activation)
        previous_layer = current_layer
    return output_convs


# 从cfg中解析出分类数
# 返回值：int 分类个数
def get_category_num(cfg_file_path):
    """Find number of output classes of the yolo model."""
    with open(cfg_file_path, 'r') as f:
        cfg_lines = [l.strip() for l in f.readlines()]
    classes_lines = [l for l in cfg_lines if l.startswith('classes=')]
    assert len(set(classes_lines)) == 1
    return int(classes_lines[-1].split('=')[-1].strip())


def get_h_and_w(layer_configs):
    """Find input height and width of the yolo model from layer configs."""
    net_config = layer_configs['000_net']
    return net_config['height'], net_config['width']


def get_anchors(cfg_file_path):
    """Get anchors of all yolo layers from the cfg file."""
    with open(cfg_file_path, 'r') as f:
        cfg_lines = f.readlines()
    yolo_lines = [l.strip() for l in cfg_lines if l.startswith('[yolo]')]
    mask_lines = [l.strip() for l in cfg_lines if l.startswith('mask')]
    anch_lines = [l.strip() for l in cfg_lines if l.startswith('anchors')]
    assert len(mask_lines) == len(yolo_lines)
    assert len(anch_lines) == len(yolo_lines)
    anchor_list = eval('[%s]' % anch_lines[0].split('=')[-1])
    mask_strs = [l.split('=')[-1] for l in mask_lines]
    masks = [eval('[%s]' % s)  for s in mask_strs]
    anchors = []
    for mask in masks:
        curr_anchors = []
        for m in mask:
            curr_anchors.append(anchor_list[m * 2])
            curr_anchors.append(anchor_list[m * 2 + 1])
        anchors.append(curr_anchors)
    return anchors


def get_anchor_num(cfg_file_path):
    """Find number of anchors (masks) of the yolo model."""
    anchors = get_anchors(cfg_file_path)
    num_anchors = [len(a) // 2 for a in anchors]

    assert len(num_anchors) > 0, 'Found no `mask` fields in config'
    assert len(set(num_anchors)) == 1, 'Found different num anchors'

    return num_anchors[0]


class DarkNetParser(object):
    """Definition of a parser for DarkNet-based YOLO model."""

    def __init__(self, supported_layers=None):
        """Initializes a DarkNetParser object.

        Keyword argument:
        supported_layers -- a string list of supported layers in DarkNet naming convention,
        parameters are only added to the class dictionary if a parsed layer is included.
        """

        # A list of YOLO layers containing dictionaries with all layer
        # parameters:

        # OrderedDic 可以保持键值对的顺序[]
        self.layer_configs = OrderedDict()
        # 支持的节点类型
        self.supported_layers = supported_layers if supported_layers else \
                                ['net',             # 超参数层，无操作
                                 'convolutional',   # 卷积层
                                 'maxpool',         # 池化层
                                 'shortcut',        # 捷径层
                                 'route',           # 路由层
                                 'upsample',        # 上采样层
                                 'yolo']            # 输出层
        self.layer_counter = 0

    # 加载网络模型文件.cfg
    def parse_cfg_file(self, cfg_file_path):
        """Takes the yolov?.cfg file and parses it layer by layer,
        appending each layer's parameters as a dictionary to layer_configs.

        Keyword argument:
        cfg_file_path
        """
        with open(cfg_file_path, 'r') as cfg_file:
            remainder = cfg_file.read()
            while remainder is not None:
                # 从字符串中加载一层网络，并生成字典结构
                layer_dict, layer_name, remainder = self._next_layer(remainder)
                if layer_dict is not None:
                    # 将一层网络结构根据名称，依次加入有序字典中，生成整个网络结构
                    self.layer_configs[layer_name] = layer_dict
        return self.layer_configs

    # 返回当前层生成的字典键值对，并指向下一层结构
    # layer_dict 一层内的网络结构字典， layer_name 层名称， remainder 剩余字符串
    def _next_layer(self, remainder):
        """Takes in a string and segments it by looking for DarkNet delimiters.
        Returns the layer parameters and the remaining string after the last delimiter.
        Example for the first Conv layer in yolo.cfg ...

        [convolutional]
        batch_normalize=1
        filters=32
        size=3
        stride=1
        pad=1
        activation=leaky

        ... becomes the following layer_dict return value:
        {'activation': 'leaky', 'stride': 1, 'pad': 1, 'filters': 32,
        'batch_normalize': 1, 'type': 'convolutional', 'size': 3}.

        '001_convolutional' is returned as layer_name, and all lines that follow in yolo.cfg
        are returned as the next remainder.

        Keyword argument:
        remainder -- a string with all raw text after the previously parsed layer
        """
        remainder = remainder.split('[', 1)
        while len(remainder[0]) > 0 and remainder[0][-1] == '#':
            # '#[...' case (the left bracket is proceeded by a pound sign),
            # assuming this layer is commented out, so go find the next '['
            remainder = remainder[1].split('[', 1)
        if len(remainder) == 2:
            remainder = remainder[1]
        else:
            # no left bracket found in remainder
            return None, None, None
        remainder = remainder.split(']', 1)
        if len(remainder) == 2:
            layer_type, remainder = remainder
        else:
            # no right bracket
            raise ValueError('no closing bracket!')
        if layer_type not in self.supported_layers:
            raise ValueError('%s layer not supported!' % layer_type)

        out = remainder.split('\n[', 1)
        if len(out) == 2:
            layer_param_block, remainder = out[0], '[' + out[1]
        else:
            layer_param_block, remainder = out[0], ''
        layer_param_lines = layer_param_block.split('\n')
        # remove empty lines
        layer_param_lines = [l.lstrip() for l in layer_param_lines if l.lstrip()]
        # don't parse yolo layers
        if layer_type == 'yolo':  layer_param_lines = []
        skip_params = ['steps', 'scales'] if layer_type == 'net' else []
        layer_name = str(self.layer_counter).zfill(3) + '_' + layer_type
        layer_dict = dict(type=layer_type)
        for param_line in layer_param_lines:
            param_line = param_line.split('#')[0]
            if not param_line:  continue
            assert '[' not in param_line
            param_type, param_value = self._parse_params(param_line, skip_params)
            layer_dict[param_type] = param_value
        self.layer_counter += 1
        return layer_dict, layer_name, remainder

    def _parse_params(self, param_line, skip_params=None):
        """Identifies the parameters contained in one of the cfg file and returns
        them in the required format for each parameter type, e.g. as a list, an int or a float.

        Keyword argument:
        param_line -- one parsed line within a layer block
        """
        param_line = param_line.replace(' ', '')
        param_type, param_value_raw = param_line.split('=')
        assert param_value_raw
        param_value = None
        if skip_params and param_type in skip_params:
            param_type = None
        elif param_type == 'layers':
            layer_indexes = list()
            for index in param_value_raw.split(','):
                layer_indexes.append(int(index))
            param_value = layer_indexes
        elif isinstance(param_value_raw, str) and not param_value_raw.isalpha():
            condition_param_value_positive = param_value_raw.isdigit()
            condition_param_value_negative = param_value_raw[0] == '-' and \
                param_value_raw[1:].isdigit()
            if condition_param_value_positive or condition_param_value_negative:
                param_value = int(param_value_raw)
            else:
                param_value = float(param_value_raw)
        else:
            param_value = str(param_value_raw)
        return param_type, param_value


class MajorNodeSpecs(object):
    """Helper class used to store the names of ONNX output names,
    corresponding to the output of a DarkNet layer and its output channels.
    Some DarkNet layers are not created and there is no corresponding ONNX node,
    but we still need to track them in order to set up skip connections.
    """

    def __init__(self, name, channels):
        """ Initialize a MajorNodeSpecs object.

        Keyword arguments:
        name -- name of the ONNX node
        channels -- number of output channels of this node
        """
        self.name = name
        self.channels = channels
        self.created_onnx_node = False
        if name is not None and isinstance(channels, int) and channels > 0:
            self.created_onnx_node = True


class ConvParams(object):
    """Helper class to store the hyper parameters of a Conv layer,
    including its prefix name in the ONNX graph and the expected dimensions
    of weights for convolution, bias, and batch normalization.

    Additionally acts as a wrapper for generating safe names for all
    weights, checking on feasible combinations.
    """

    def __init__(self, node_name, batch_normalize, conv_weight_dims):
        """Constructor based on the base node name (e.g. 101_convolutional), the batch
        normalization setting, and the convolutional weights shape.

        Keyword arguments:
        node_name -- base name of this YOLO convolutional layer
        batch_normalize -- bool value if batch normalization is used
        conv_weight_dims -- the dimensions of this layer's convolutional weights
        """
        self.node_name = node_name
        self.batch_normalize = batch_normalize
        assert len(conv_weight_dims) == 4
        self.conv_weight_dims = conv_weight_dims

    def generate_param_name(self, param_category, suffix):
        """Generates a name based on two string inputs,
        and checks if the combination is valid."""
        assert suffix
        assert param_category in ['bn', 'conv']
        assert(suffix in ['scale', 'mean', 'var', 'weights', 'bias'])
        if param_category == 'bn':
            assert self.batch_normalize
            assert suffix in ['scale', 'bias', 'mean', 'var']
        elif param_category == 'conv':
            assert suffix in ['weights', 'bias']
            if suffix == 'bias':
                assert not self.batch_normalize
        param_name = self.node_name + '_' + param_category + '_' + suffix
        return param_name

class ResizeParams(object):
    #Helper class to store the scale parameter for an Resize node.

    def __init__(self, node_name, value):
        """Constructor based on the base node name (e.g. 86_Resize),
        and the value of the scale input tensor.

        Keyword arguments:
        node_name -- base name of this YOLO Resize layer
        value -- the value of the scale input to the Resize layer as numpy array
        """
        self.node_name = node_name
        self.value = value

    def generate_param_name(self):
        """Generates the scale parameter name for the Resize node."""
        param_name = self.node_name + '_' + "scale"
        return param_name

    def generate_roi_name(self):
        """Generates the roi input name for the Resize node."""
        param_name = self.node_name + '_' + "roi"
        return param_name

class WeightLoader(object):
    """Helper class used for loading the serialized weights of a binary file stream
    and returning the initializers and the input tensors required for populating
    the ONNX graph with weights.
    """

    def __init__(self, weights_file_path):
        """Initialized with a path to the YOLO .weights file.

        Keyword argument:
        weights_file_path -- path to the weights file.
        """
        self.weights_file = self._open_weights_file(weights_file_path)

    def load_resize_scales(self, resize_params):
        """Returns the initializers with the value of the scale input
        tensor given by resize_params.

        Keyword argument:
        resize_params -- a ResizeParams object
        """
        initializer = list()
        inputs = list()
        name = resize_params.generate_param_name()
        shape = resize_params.value.shape
        data = resize_params.value
        scale_init = helper.make_tensor(
            name, TensorProto.FLOAT, shape, data)
        scale_input = helper.make_tensor_value_info(
            name, TensorProto.FLOAT, shape)
        initializer.append(scale_init)
        inputs.append(scale_input)

        # In opset 11 an additional input named roi is required. Create a dummy tensor to satisfy this.
        # It is a 1D tensor of size of the rank of the input (4)
        rank = 4
        roi_name = resize_params.generate_roi_name()
        roi_input = helper.make_tensor_value_info(roi_name, TensorProto.FLOAT, [rank])
        roi_init = helper.make_tensor(roi_name, TensorProto.FLOAT, [rank], [0,0,0,0])
        initializer.append(roi_init)
        inputs.append(roi_input)

        return initializer, inputs

    def load_conv_weights(self, conv_params):
        """Returns the initializers with weights from the weights file and
        the input tensors of a convolutional layer for all corresponding ONNX nodes.

        Keyword argument:
        conv_params -- a ConvParams object
        """
        initializer = list()
        inputs = list()
        if conv_params.batch_normalize:
            bias_init, bias_input = self._create_param_tensors(
                conv_params, 'bn', 'bias')
            bn_scale_init, bn_scale_input = self._create_param_tensors(
                conv_params, 'bn', 'scale')
            bn_mean_init, bn_mean_input = self._create_param_tensors(
                conv_params, 'bn', 'mean')
            bn_var_init, bn_var_input = self._create_param_tensors(
                conv_params, 'bn', 'var')
            initializer.extend(
                [bn_scale_init, bias_init, bn_mean_init, bn_var_init])
            inputs.extend([bn_scale_input, bias_input,
                           bn_mean_input, bn_var_input])
        else:
            bias_init, bias_input = self._create_param_tensors(
                conv_params, 'conv', 'bias')
            initializer.append(bias_init)
            inputs.append(bias_input)
        conv_init, conv_input = self._create_param_tensors(
            conv_params, 'conv', 'weights')
        initializer.append(conv_init)
        inputs.append(conv_input)
        return initializer, inputs

    def _open_weights_file(self, weights_file_path):
        """Opens a YOLO DarkNet file stream and skips the header.

        Keyword argument:
        weights_file_path -- path to the weights file.
        """
        weights_file = open(weights_file_path, 'rb')
        length_header = 5
        np.ndarray(shape=(length_header, ), dtype='int32',
                   buffer=weights_file.read(length_header * 4))
        return weights_file

    def _create_param_tensors(self, conv_params, param_category, suffix):
        """Creates the initializers with weights from the weights file together with
        the input tensors.

        Keyword arguments:
        conv_params -- a ConvParams object
        param_category -- the category of parameters to be created ('bn' or 'conv')
        suffix -- a string determining the sub-type of above param_category (e.g.,
        'weights' or 'bias')
        """
        param_name, param_data, param_data_shape = self._load_one_param_type(
            conv_params, param_category, suffix)

        initializer_tensor = helper.make_tensor(
            param_name, TensorProto.FLOAT, param_data_shape, param_data)
        input_tensor = helper.make_tensor_value_info(
            param_name, TensorProto.FLOAT, param_data_shape)
        return initializer_tensor, input_tensor

    def _load_one_param_type(self, conv_params, param_category, suffix):
        """Deserializes the weights from a file stream in the DarkNet order.

        Keyword arguments:
        conv_params -- a ConvParams object
        param_category -- the category of parameters to be created ('bn' or 'conv')
        suffix -- a string determining the sub-type of above param_category (e.g.,
        'weights' or 'bias')
        """
        param_name = conv_params.generate_param_name(param_category, suffix)
        channels_out, channels_in, filter_h, filter_w = conv_params.conv_weight_dims
        if param_category == 'bn':
            param_shape = [channels_out]
        elif param_category == 'conv':
            if suffix == 'weights':
                param_shape = [channels_out, channels_in, filter_h, filter_w]
            elif suffix == 'bias':
                param_shape = [channels_out]
        param_size = np.product(np.array(param_shape))
        param_data = np.ndarray(
            shape=param_shape,
            dtype='float32',
            buffer=self.weights_file.read(param_size * 4))
        param_data = param_data.flatten().astype(float)
        return param_name, param_data, param_shape


class GraphBuilderONNX(object):
    """Class for creating an ONNX graph from a previously generated list of layer dictionaries."""

    def __init__(self, model_name, output_tensors, batch_size):
        """Initialize with all DarkNet default parameters used creating
        YOLO, and specify the output tensors as an OrderedDict for their
        output dimensions with their names as keys.

        Keyword argument:
        output_tensors -- the output tensors as an OrderedDict containing the keys'
        output dimensions
        """
        self.model_name = model_name
        self.output_tensors = output_tensors
        self._nodes = list()
        self.graph_def = None
        self.input_tensor = None
        self.epsilon_bn = 1e-5
        self.momentum_bn = 0.99
        self.alpha_lrelu = 0.1
        self.param_dict = OrderedDict()
        self.major_node_specs = list()
        self.batch_size = batch_size
        self.route_spec = 0  # keeping track of the current active 'route'

    def build_onnx_graph(
            self,
            layer_configs,
            weights_file_path,
            verbose=True):
        """Iterate over all layer configs (parsed from the DarkNet
        representation of YOLO), create an ONNX graph, populate it with
        weights from the weights file and return the graph definition.

        Keyword arguments:
        layer_configs -- an OrderedDict object with all parsed layers' configurations
        weights_file_path -- location of the weights file
        verbose -- toggles if the graph is printed after creation (default: True)
        """
        for layer_name in layer_configs.keys():
            layer_dict = layer_configs[layer_name]
            # 根据网络结构，分别生成onnx节点，节点的作用是操作符，如conv,relu等，每个操作都要写成Onnx格式
            major_node_specs = self._make_onnx_node(layer_name, layer_dict)
            # 成功生成后，添加到主网络结构中
            if major_node_specs.name is not None:
                self.major_node_specs.append(major_node_specs)
        # remove dummy 'route' and 'yolo' nodes
        self.major_node_specs = [node for node in self.major_node_specs
                                      if 'dummy' not in node.name]
        outputs = list()
        # 遍历输出字典中的名称
        for tensor_name in self.output_tensors.keys():
            # 输出维度，例[batch,255,13,13]
            output_dims = [self.batch_size, ] + \
                self.output_tensors[tensor_name]
            # 创建Onnx的"变量"（ValueInfoProto）
            output_tensor = helper.make_tensor_value_info(
                tensor_name, TensorProto.FLOAT, output_dims)
            # 添加到输出列表中
            outputs.append(output_tensor)
        inputs = [self.input_tensor]
        # 加载权重到ndarray中，weight文件如何存储的？？按顺序存储的二进制文件
        weight_loader = WeightLoader(weights_file_path)
        initializer = list()
        # If a layer has parameters, add them to the initializer and input lists. ？？？
        # 大概是按照layer生成各层级的节点信息，并保存权重（darknet格式）到节点中
        # initializer是包含权重信息的，input则是ValueInfoProto，可以理解为变量
        for layer_name in self.param_dict.keys():
            _, layer_type = layer_name.split('_', 1)
            params = self.param_dict[layer_name]
            if layer_type == 'convolutional':
                initializer_layer, inputs_layer = weight_loader.load_conv_weights(
                    params)
                initializer.extend(initializer_layer)
                inputs.extend(inputs_layer)
            elif layer_type == 'upsample':
                initializer_layer, inputs_layer = weight_loader.load_resize_scales(
                    params)
                initializer.extend(initializer_layer)
                inputs.extend(inputs_layer)
        del weight_loader
        # 生成onnx图
        self.graph_def = helper.make_graph(
            nodes=self._nodes,
            name=self.model_name,
            inputs=inputs,
            outputs=outputs,
            initializer=initializer
        )
        if verbose:
            print(helper.printable_graph(self.graph_def))
        model_def = helper.make_model(self.graph_def,
                                      producer_name='NVIDIA TensorRT sample')
        return model_def

    def _make_onnx_node(self, layer_name, layer_dict):
        """Take in a layer parameter dictionary, choose the correct function for
        creating an ONNX node and store the information important to graph creation
        as a MajorNodeSpec object.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        layer_type = layer_dict['type']
        if self.input_tensor is None:
            if layer_type == 'net':
                major_node_output_name, major_node_output_channels = self._make_input_tensor(
                    layer_name, layer_dict)
                major_node_specs = MajorNodeSpecs(major_node_output_name,
                                                  major_node_output_channels)
            else:
                raise ValueError('The first node has to be of type "net".')
        else:
            node_creators = dict()
            node_creators['convolutional'] = self._make_conv_node
            node_creators['maxpool'] = self._make_maxpool_node
            node_creators['shortcut'] = self._make_shortcut_node
            node_creators['route'] = self._make_route_node
            node_creators['upsample'] = self._make_resize_node
            node_creators['yolo'] = self._make_yolo_node

            if layer_type in node_creators.keys():
                major_node_output_name, major_node_output_channels = \
                    node_creators[layer_type](layer_name, layer_dict)
                major_node_specs = MajorNodeSpecs(major_node_output_name,
                                                  major_node_output_channels)
            else:
                raise TypeError('layer of type %s not supported' % layer_type)
        return major_node_specs

    def _make_input_tensor(self, layer_name, layer_dict):
        """Create an ONNX input tensor from a 'net' layer and store the batch size.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        #batch_size = layer_dict['batch']
        channels = layer_dict['channels']
        height = layer_dict['height']
        width = layer_dict['width']
        #self.batch_size = batch_size
        input_tensor = helper.make_tensor_value_info(
            str(layer_name), TensorProto.FLOAT, [
                self.batch_size, channels, height, width])
        self.input_tensor = input_tensor
        return layer_name, channels

    def _get_previous_node_specs(self, target_index=0):
        """Get a previously ONNX node.

        Target index can be passed for jumping to a specific index.

        Keyword arguments:
        target_index -- optional for jumping to a specific index,
                        default: 0 for the previous element, while
                        taking 'route' spec into account
        """
        if target_index == 0:
            if self.route_spec != 0:
                previous_node = self.major_node_specs[self.route_spec]
                assert 'dummy' not in previous_node.name
                self.route_spec = 0
            else:
                previous_node = self.major_node_specs[-1]
        else:
            previous_node = self.major_node_specs[target_index]
        assert previous_node.created_onnx_node
        return previous_node

    def _make_conv_node(self, layer_name, layer_dict):
        """Create an ONNX Conv node with optional batch normalization and
        activation nodes.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        previous_node_specs = self._get_previous_node_specs()
        inputs = [previous_node_specs.name]
        previous_channels = previous_node_specs.channels
        kernel_size = layer_dict['size']
        stride = layer_dict['stride']
        filters = layer_dict['filters']
        batch_normalize = False
        if layer_dict.get('batch_normalize', 0) > 0:
            batch_normalize = True

        kernel_shape = [kernel_size, kernel_size]
        weights_shape = [filters, previous_channels] + kernel_shape
        conv_params = ConvParams(layer_name, batch_normalize, weights_shape)

        strides = [stride, stride]
        dilations = [1, 1]
        weights_name = conv_params.generate_param_name('conv', 'weights')
        inputs.append(weights_name)
        if not batch_normalize:
            bias_name = conv_params.generate_param_name('conv', 'bias')
            inputs.append(bias_name)

        conv_node = helper.make_node(
            'Conv',
            inputs=inputs,
            outputs=[layer_name],
            kernel_shape=kernel_shape,
            strides=strides,
            auto_pad='SAME_LOWER',
            dilations=dilations,
            name=layer_name
        )
        self._nodes.append(conv_node)
        inputs = [layer_name]
        layer_name_output = layer_name

        if batch_normalize:
            layer_name_bn = layer_name + '_bn'
            bn_param_suffixes = ['scale', 'bias', 'mean', 'var']
            for suffix in bn_param_suffixes:
                bn_param_name = conv_params.generate_param_name('bn', suffix)
                inputs.append(bn_param_name)
            batchnorm_node = helper.make_node(
                'BatchNormalization',
                inputs=inputs,
                outputs=[layer_name_bn],
                epsilon=self.epsilon_bn,
                momentum=self.momentum_bn,
                name=layer_name_bn
            )
            self._nodes.append(batchnorm_node)
            inputs = [layer_name_bn]
            layer_name_output = layer_name_bn

        if layer_dict['activation'] == 'leaky':
            layer_name_lrelu = layer_name + '_lrelu'

            lrelu_node = helper.make_node(
                'LeakyRelu',
                inputs=inputs,
                outputs=[layer_name_lrelu],
                name=layer_name_lrelu,
                alpha=self.alpha_lrelu
            )
            self._nodes.append(lrelu_node)
            inputs = [layer_name_lrelu]
            layer_name_output = layer_name_lrelu
        elif layer_dict['activation'] == 'mish':
            layer_name_softplus = layer_name + '_softplus'
            layer_name_tanh = layer_name + '_tanh'
            layer_name_mish = layer_name + '_mish'

            softplus_node = helper.make_node(
                'Softplus',
                inputs=inputs,
                outputs=[layer_name_softplus],
                name=layer_name_softplus
            )
            self._nodes.append(softplus_node)
            tanh_node = helper.make_node(
                'Tanh',
                inputs=[layer_name_softplus],
                outputs=[layer_name_tanh],
                name=layer_name_tanh
            )
            self._nodes.append(tanh_node)

            inputs.append(layer_name_tanh)
            mish_node = helper.make_node(
                'Mul',
                inputs=inputs,
                outputs=[layer_name_mish],
                name=layer_name_mish
            )
            self._nodes.append(mish_node)

            inputs = [layer_name_mish]
            layer_name_output = layer_name_mish
        elif layer_dict['activation'] == 'swish':
            layer_name_sigmoid = layer_name + '_sigmoid'
            layer_name_swish = layer_name + '_swish'

            sigmoid_node = helper.make_node(
                'Sigmoid',
                inputs=inputs,
                outputs=[layer_name_sigmoid],
                name=layer_name_sigmoid
            )
            self._nodes.append(sigmoid_node)

            inputs.append(layer_name_sigmoid)
            swish_node = helper.make_node(
                'Mul',
                inputs=inputs,
                outputs=[layer_name_swish],
                name=layer_name_swish
            )
            self._nodes.append(swish_node)

            inputs = [layer_name_swish]
            layer_name_output = layer_name_swish
        elif layer_dict['activation'] == 'logistic':
            layer_name_lgx = layer_name + '_lgx'

            lgx_node = helper.make_node(
                'Sigmoid',
                inputs=inputs,
                outputs=[layer_name_lgx],
                name=layer_name_lgx
            )
            self._nodes.append(lgx_node)
            inputs = [layer_name_lgx]
            layer_name_output = layer_name_lgx
        elif layer_dict['activation'] == 'linear':
            pass
        else:
            raise TypeError('%s activation not supported' % layer_dict['activation'])

        self.param_dict[layer_name] = conv_params
        return layer_name_output, filters

    def _make_shortcut_node(self, layer_name, layer_dict):
        """Create an ONNX Add node with the shortcut properties from
        the DarkNet-based graph.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        shortcut_index = layer_dict['from']
        activation = layer_dict['activation']
        assert activation == 'linear'

        first_node_specs = self._get_previous_node_specs()
        second_node_specs = self._get_previous_node_specs(
            target_index=shortcut_index)
        assert first_node_specs.channels == second_node_specs.channels
        channels = first_node_specs.channels
        inputs = [first_node_specs.name, second_node_specs.name]
        shortcut_node = helper.make_node(
            'Add',
            inputs=inputs,
            outputs=[layer_name],
            name=layer_name,
        )
        self._nodes.append(shortcut_node)
        return layer_name, channels

    def _make_route_node(self, layer_name, layer_dict):
        """If the 'layers' parameter from the DarkNet configuration is only one index, continue
        node creation at the indicated (negative) index. Otherwise, create an ONNX Concat node
        with the route properties from the DarkNet-based graph.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        route_node_indexes = layer_dict['layers']
        if len(route_node_indexes) == 1:
            if 'groups' in layer_dict.keys():
                # for CSPNet-kind of architecture
                assert 'group_id' in layer_dict.keys()
                groups = layer_dict['groups']
                group_id = int(layer_dict['group_id'])
                assert group_id < groups
                index = route_node_indexes[0]
                if index > 0:
                    # +1 for input node (same reason as below)
                    index += 1
                route_node_specs = self._get_previous_node_specs(
                    target_index=index)
                assert route_node_specs.channels % groups == 0
                channels = route_node_specs.channels // groups

                outputs = [layer_name + '_dummy%d' % i for i in range(groups)]
                outputs[group_id] = layer_name
                route_node = helper.make_node(
                    'Split',
                    axis=1,
                    #split=[channels] * groups,  # not needed for opset 11
                    inputs=[route_node_specs.name],
                    outputs=outputs,
                    name=layer_name,
                )
                self._nodes.append(route_node)
            else:
                if route_node_indexes[0] < 0:
                    # route should skip self, thus -1
                    self.route_spec = route_node_indexes[0] - 1
                elif route_node_indexes[0] > 0:
                    # +1 for input node (same reason as below)
                    self.route_spec = route_node_indexes[0] + 1
                # This dummy route node would be removed in the end.
                layer_name = layer_name + '_dummy'
                channels = 1
        else:
            assert 'groups' not in layer_dict.keys(), \
                'groups not implemented for multiple-input route layer!'
            inputs = list()
            channels = 0
            for index in route_node_indexes:
                if index > 0:
                    # Increment by one because we count the input as
                    # a node (DarkNet does not)
                    index += 1
                route_node_specs = self._get_previous_node_specs(
                    target_index=index)
                inputs.append(route_node_specs.name)
                channels += route_node_specs.channels
            assert inputs
            assert channels > 0

            route_node = helper.make_node(
                'Concat',
                axis=1,
                inputs=inputs,
                outputs=[layer_name],
                name=layer_name,
            )
            self._nodes.append(route_node)
        return layer_name, channels

    def _make_resize_node(self, layer_name, layer_dict):
        """Create an ONNX Resize node with the properties from
        the DarkNet-based graph.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        resize_scale_factors = float(layer_dict['stride'])
        # Create the scale factor array with node parameters
        scales=np.array([1.0, 1.0, resize_scale_factors, resize_scale_factors]).astype(np.float32)
        previous_node_specs = self._get_previous_node_specs()
        inputs = [previous_node_specs.name]

        channels = previous_node_specs.channels
        assert channels > 0
        resize_params = ResizeParams(layer_name, scales)

        # roi input is the second input, so append it before scales
        roi_name = resize_params.generate_roi_name()
        inputs.append(roi_name)

        scales_name = resize_params.generate_param_name()
        inputs.append(scales_name)

        resize_node = helper.make_node(
            'Resize',
            coordinate_transformation_mode='asymmetric',
            mode='nearest',
            nearest_mode='floor',
            inputs=inputs,
            outputs=[layer_name],
            name=layer_name,
        )
        self._nodes.append(resize_node)
        self.param_dict[layer_name] = resize_params
        return layer_name, channels

    def _make_maxpool_node(self, layer_name, layer_dict):
        """Create an ONNX Maxpool node with the properties from
        the DarkNet-based graph.

        Keyword arguments:
        layer_name -- the layer's name (also the corresponding key in layer_configs)
        layer_dict -- a layer parameter dictionary (one element of layer_configs)
        """
        stride = layer_dict['stride']
        kernel_size = layer_dict['size']
        previous_node_specs = self._get_previous_node_specs()
        inputs = [previous_node_specs.name]
        channels = previous_node_specs.channels
        kernel_shape = [kernel_size, kernel_size]
        strides = [stride, stride]
        assert channels > 0
        maxpool_node = helper.make_node(
            'MaxPool',
            inputs=inputs,
            outputs=[layer_name],
            kernel_shape=kernel_shape,
            strides=strides,
            auto_pad='SAME_UPPER',
            name=layer_name,
        )
        self._nodes.append(maxpool_node)
        return layer_name, channels

    def _make_yolo_node(self, layer_name, layer_dict):
        """Create an ONNX Yolo node.

        These are dummy nodes which would be removed in the end.
        """
        channels = 1
        return layer_name + '_dummy', channels


def main():
    if sys.version_info[0] < 3:
        raise SystemExit('ERROR: This modified version of yolov3_to_onnx.py '
                         'script is only compatible with python3...')

    args = parse_args()
    # 网络模型路径
    cfg_file_path = '%s.cfg' % args.model
    if not os.path.isfile(cfg_file_path):
        raise SystemExit('ERROR: file (%s) not found!' % cfg_file_path)
    # 权重模型路径
    weights_file_path = '%s.weights' % args.model
    if not os.path.isfile(weights_file_path):
        raise SystemExit('ERROR: file (%s) not found!' % weights_file_path)
    output_file_path = '%s.onnx' % args.model

    # Darknet模型解释器(.cfg->.onnx格式)
    print('Parsing DarkNet cfg file...')
    parser = DarkNetParser()
    # 从cfg文件中加载网络结构，并按顺序存入字典中[layer_name,[key,value]]
    layer_configs = parser.parse_cfg_file(cfg_file_path)
    # 获取yolo输出分类个数，单位int
    category_num = get_category_num(cfg_file_path)
    # 获取输出层名称从yolo层提取，用于推算后获取结果
    output_tensor_names = get_output_convs(layer_configs)
    # e.g. ['036_convolutional', '044_convolutional', '052_convolutional']

    # 获取输出维度 (80 + 5) * 3 = 255
    c = (category_num + 5) * get_anchor_num(cfg_file_path)
    # 获取输入图像宽高
    h, w = get_h_and_w(layer_configs)
    # 获取输出格式
    if len(output_tensor_names) == 2:
        # 2种候选框
        output_tensor_shapes = [
            [c, h // 32, w // 32], [c, h // 16, w // 16]]
    elif len(output_tensor_names) == 3:
        # 3种候选框
        output_tensor_shapes = [
            [c, h // 32, w // 32], [c, h // 16, w // 16],
            [c, h // 8, w // 8]]
    elif len(output_tensor_names) == 4:
        # 4种候选框
        output_tensor_shapes = [
            [c, h // 64, w // 64], [c, h // 32, w // 32],
            [c, h // 16, w // 16], [c, h // 8, w // 8]]
    # 判断是金字塔模式（自下而上）还是倒金字塔模式（自上而下），决定了输出的顺序
    if is_pan_arch(cfg_file_path):
        # 输出从大图到小图，改为小图先输出
        output_tensor_shapes.reverse()
    # 生成输出字典格式,以416为例 [036_convolutional，255，13，13],[044_convolutional,255,26,26],[052_convolutional,255,52,52]
    output_tensor_dims = OrderedDict(
        zip(output_tensor_names, output_tensor_shapes))

    # 创建ONNX生成器
    print('Building ONNX graph...')
    builder = GraphBuilderONNX(
        args.model,                 # Pytorch模型
        output_tensor_dims,         # 输出字典[[],[],[]]
        MAX_BATCH_SIZE)             # 最大Batch数量
    # 编译ONNX模型
    yolo_model_def = builder.build_onnx_graph(
        layer_configs=layer_configs,            # 网络层结构
        weights_file_path=weights_file_path,    # 网络权重
        verbose=True)                           # 显示生成过程

    # 检查生成的ONNX模型
    print('Checking ONNX model...')
    onnx.checker.check_model(yolo_model_def)

    # 保存ONNX模型
    print('Saving ONNX file...')
    onnx.save(yolo_model_def, output_file_path)

    print('Done.')


if __name__ == '__main__':
    main()