1107 lines
45 KiB
Python
1107 lines
45 KiB
Python
# 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()
|