原官方keras 版 maskrcnn 转onnx，并使用onnxruntime gpu c++ 推理
Posted 东东就是我
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了原官方keras 版 maskrcnn 转onnx，并使用onnxruntime gpu c++ 推理相关的知识，希望对你有一定的参考价值。
https://github.com/dulvqingyun/onnxruntime-maskrcnn/blob/master/newC%2B%2B
1.keras 下模型转onnx
import os
import sys
import numpy as np
import skimage
import onnx
import keras2onnx
from mrcnn.config import Config
from mrcnn.model import BatchNorm, DetectionLayer
from mrcnn import model as modellib
from mrcnn import visualize
from keras2onnx import set_converter
from keras2onnx.ke2onnx.batch_norm import convert_keras_batch_normalization
from keras2onnx.proto import onnx_proto
from keras2onnx.common.onnx_ops import apply_transpose, apply_identity
from keras2onnx.common.onnx_ops import OnnxOperatorBuilder
from os.path import dirname, abspath
sys.path.insert(0, os.path.join(dirname(abspath(__file__)), '../../tests/'))
# from test_utils import convert_tf_crop_and_resize
ROOT_DIR = os.path.abspath("./")
# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
# Path to trained weights file
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
class CocoConfig(Config):
"""Configuration for training on MS COCO.
Derives from the base Config class and overrides values specific
to the COCO dataset.
"""
# Give the configuration a recognizable name
NAME = "coco"
# We use a GPU with 12GB memory, which can fit two images.
# Adjust down if you use a smaller GPU.
IMAGES_PER_GPU = 2
# Uncomment to train on 8 GPUs (default is 1)
# GPU_COUNT = 8
# Number of classes (including background)
NUM_CLASSES = 1 + 1 # COCO has 80 classes
class InferenceConfig(CocoConfig):
# Set batch size to 1 since we'll be running inference on
# one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
GPU_COUNT = 1
IMAGES_PER_GPU = 1
config = InferenceConfig()
config.display()
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
# Load weights trained on MS-COCO
model.load_weights('samples/balloon/mask_rcnn_balloon.h5', by_name=True)
def convert_tf_crop_and_resize(scope, operator, container):
if operator.target_opset < 11:
raise ValueError("CropAndResize op is not supported for opset < 11")
oopb = OnnxOperatorBuilder(container, scope)
node = operator.raw_operator
mode_value = node.get_attr('method')
transpose_node = oopb.apply_transpose(operator.inputs[0].full_name,
name=operator.full_name + '_transpose_1',
perm=[0, 3, 1, 2])
cropandresize = oopb.add_node('CropAndResize',
transpose_node + operator.input_full_names[1:],
operator.full_name + '_crop_and_resize',
op_domain='com.microsoft',
op_version=1,
mode=mode_value)
oopb.apply_op_with_output("apply_transpose",
cropandresize,
operator.output_full_names,
name=operator.full_name + '_transpose_final',
perm=[0, 2, 3, 1])
def convert_BatchNorm(scope, operator, container):
convert_keras_batch_normalization(scope, operator, container)
def convert_apply_box_deltas_graph(scope, operator, container, oopb, box_transpose, score_identity, deltas_transpose, windows_transpose):
oopb = OnnxOperatorBuilder(container, scope)
box_squeeze = oopb.apply_squeeze(box_transpose, name=operator.full_name + '_box_squeeze', axes=[0])[0]
# output shape: [spatial_dimension, 4]
deltas_squeeze = oopb.apply_squeeze(deltas_transpose, name=operator.full_name + '_deltas_squeeze', axes=[0])[0]
# output shape: [spatial_dimension, num_classes, 4]
score_squeeze = oopb.apply_squeeze(score_identity, name=operator.full_name + '_score_squeeze', axes=[0])[0]
# output shape: [spatial_dimension, num_classes]
class_ids = scope.get_unique_variable_name('class_ids')
attrs = 'axis': 1
container.add_node('ArgMax', score_squeeze, class_ids, op_version=operator.target_opset,
**attrs)
# output shape: [spatial_dimension, 1]
prob_shape = oopb.add_node('Shape',
[score_squeeze],
operator.inputs[1].full_name + '_prob_shape')
prob_shape_0 = oopb.add_node('Slice',
[prob_shape,
('_start', oopb.int64, np.array([0], dtype='int64')),
('_end', oopb.int64, np.array([1], dtype='int64')),
('_axes', oopb.int64, np.array([0], dtype='int64'))
],
operator.inputs[1].full_name + '_prob_shape_0')
prob_range = oopb.add_node('Range',
[('_start', oopb.int64, np.array([0], dtype='int64')),
prob_shape_0,
# ('_limit', oopb.int64, np.array([1000], dtype='int64')),
('_delta', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[1].full_name + '_prob_range',
op_domain='com.microsoft',
op_version=1)
prob_range_unsqueeze = oopb.apply_unsqueeze([prob_range],
operator.inputs[1].full_name + '_prob_range_unsqueeze',
axes=[1])[0]
# output shape: [spatial_dimension, 1]
attrs = 'axis': 1
indices = oopb.add_node('Concat',
[prob_range_unsqueeze,
class_ids
],
operator.inputs[1].full_name + '_indices', **attrs)
# output shape: [spatial_dimension, 2]
deltas_specific = oopb.add_node('GatherND',
[deltas_squeeze, indices],
operator.inputs[2].full_name + '_deltas_specific')
# output shape: [spatial_dimension, 4]
BBOX_STD_DEV = np.array([0.1, 0.1, 0.2, 0.2], dtype='float32')
delta_mul_output = oopb.add_node('Mul',
[deltas_specific,
('_mul_constant', oopb.float, BBOX_STD_DEV)
],
operator.inputs[2].full_name + '_mul')
# output shape: [spatial_dimension, 4]
box_0 = oopb.add_node('Slice',
[box_squeeze,
('_start', oopb.int64, np.array([0], dtype='int64')),
('_end', oopb.int64, np.array([1], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[0].full_name + '_sliced_0')
box_1 = oopb.add_node('Slice',
[box_squeeze,
('_start', oopb.int64, np.array([1], dtype='int64')),
('_end', oopb.int64, np.array([2], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[0].full_name + '_sliced_1')
box_2 = oopb.add_node('Slice',
[box_squeeze,
('_start', oopb.int64, np.array([2], dtype='int64')),
('_end', oopb.int64, np.array([3], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[0].full_name + '_sliced_2')
box_3 = oopb.add_node('Slice',
[box_squeeze,
('_start', oopb.int64, np.array([3], dtype='int64')),
('_end', oopb.int64, np.array([4], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[0].full_name + '_sliced_3')
delta_0 = oopb.add_node('Slice',
[delta_mul_output,
('_start', oopb.int64, np.array([0], dtype='int64')),
('_end', oopb.int64, np.array([1], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[3].full_name + '_sliced_0')
delta_1 = oopb.add_node('Slice',
[delta_mul_output,
('_start', oopb.int64, np.array([1], dtype='int64')),
('_end', oopb.int64, np.array([2], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[3].full_name + '_sliced_1')
delta_2 = oopb.add_node('Slice',
[delta_mul_output,
('_start', oopb.int64, np.array([2], dtype='int64')),
('_end', oopb.int64, np.array([3], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[3].full_name + '_sliced_2')
delta_3 = oopb.add_node('Slice',
[delta_mul_output,
('_start', oopb.int64, np.array([3], dtype='int64')),
('_end', oopb.int64, np.array([4], dtype='int64')),
('_axes', oopb.int64, np.array([1], dtype='int64'))
],
operator.inputs[3].full_name + '_sliced_3')
height = oopb.add_node('Sub',
[box_2, box_0],
operator.inputs[0].full_name + '_height')
width = oopb.add_node('Sub',
[box_3, box_1],
operator.inputs[0].full_name + '_width')
half_height_0 = oopb.add_node('Mul',
[height,
('_mul_constant', oopb.float, np.array([0.5], dtype='float32'))
],
operator.inputs[0].full_name + '_half_height_0')
half_width_0 = oopb.add_node('Mul',
[width,
('_mul_constant', oopb.float, np.array([0.5], dtype='float32'))
],
operator.inputs[0].full_name + '_half_width_0')
center_y_0 = oopb.add_node('Add',
[box_0, half_height_0],
operator.inputs[0].full_name + '_center_y_0')
center_x_0 = oopb.add_node('Add',
[box_1, half_width_0],
operator.inputs[0].full_name + '_center_x_0')
delta_height = oopb.add_node('Mul',
[delta_0, height],
operator.inputs[0].full_name + '_delta_height')
delta_width = oopb.add_node('Mul',
[delta_1, width],
operator.inputs[0].full_name + '_delta_width')
center_y_1 = oopb.add_node('Add',
[center_y_0, delta_height],
operator.inputs[0].full_name + '_center_y_1')
center_x_1 = oopb.add_node('Add',
[center_x_0, delta_width],
operator.inputs[0].full_name + '_center_x_1')
delta_2_exp = oopb.add_node('Exp',
[delta_2],
operator.inputs[0].full_name + '_delta_2_exp')
delta_3_exp = oopb.add_node('Exp',
[delta_3],
operator.inputs[0].full_name + '_delta_3_exp')
height_exp = oopb.add_node('Mul',
[height, delta_2_exp],
operator.inputs[0].full_name + '_height_exp')
width_exp = oopb.add_node('Mul',
[width, delta_3_exp],
operator.inputs[0].full_name + '_width_exp'以上是关于原官方keras 版 maskrcnn 转onnx，并使用onnxruntime gpu c++ 推理的主要内容，如果未能解决你的问题，请参考以下文章
『深度应用』一小时教你上手MaskRCNN·Keras开源实战（Windows&Linux）