convert model of Yolo-Fastest to onnx
why
The model generated by Yolo-Fastest cannot be converted to onnx, because grouped convolution is used many places by Yolo-Fastest
and is not supported by the scriptyolo_to_onnx.py
.
how
To support grouped convolution, just modify 3 places.
class ConvParams
def __init__(self, node_name, batch_normalize, conv_weight_dims, groups):
"""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.groups = groups
self.node_name = node_name
self.batch_normalize = batch_normalize
assert len(conv_weight_dims) == 4
self.conv_weight_dims = conv_weight_dims
self.groups = groups
function _load_one_param_type of class WeightLoader:
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))
if conv_params.groups > 1 and suffix == 'weights':
param_size = param_size // conv_params.groups
param_shape = [channels_out, channels_in//conv_params.groups, filter_h, filter_w]
buffer=self.weights_file.read(param_size * 4)
#print(param_name,param_shape,param_size*4,len(buffer))
param_data = np.ndarray(
shape=[param_size],
dtype='float32',
buffer=buffer)
param_data = param_data.flatten().astype(float)
return param_name, param_data, param_shape
function _make_conv_node of class GraphBuilderONNX
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)
"""
#print(layer_name, layer_dict)
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']
groups = layer_dict['groups'] if 'groups' in layer_dict.keys() else 1
if groups < 1:
groups = 1
batch_normalize = False
if 'batch_normalize' in layer_dict.keys(
) and layer_dict['batch_normalize'] == 1:
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, groups)
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)
padding = (kernel_size-1)//2
conv_node = helper.make_node(
'Conv',
inputs=inputs,
outputs=[layer_name],
kernel_shape=kernel_shape,
strides=strides,
pads=[padding,padding,padding+ (1 if (kernel_size-1)%2 else 0),padding+ (1 if (kernel_size-1)%2 else 0)],
#auto_pad='SAME_LOWER',
dilations=dilations,
groups=groups,
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_mish = layer_name + '_mish'
lrelu_node = helper.make_node(
'Mish',
inputs=inputs,
outputs=[layer_name_mish],
name=layer_name_mish,
)
self._nodes.append(lrelu_node)
inputs = [layer_name_mish]
layer_name_output = layer_name_mish
elif layer_dict['activation'] == 'linear':
pass
else:
print('Activation not supported.')
self.param_dict[layer_name] = conv_params
return layer_name_output, filters
Full code can be get from here.
references
- https://github.com/CaoWGG/TensorRT-YOLOv4/blob/master/tools/yolo_to_onnx.py
- https://github.com/jkjung-avt/tensorrt_demos/blob/master/yolo/yolo_to_onnx.py