def yolo2onnx(model_file, category_num=80):
"""Run the DarkNet-to-ONNX conversion for YOLO (v3 or v4)."""
if sys.version_info[0] < 3:
raise SystemExit('ERROR: This modified version of yolov3_to_onnx.py '
'script is only compatible with python3...')
parser = argparse.ArgumentParser()
parser.add_argument('-c',
'--category_num',
type=int,
default=80,
help='number of object categories [80]')
parser.add_argument(
'-m',
'--model',
type=str,
required=True,
help=('[yolov3|yolov3-tiny|yolov3-spp|yolov4|yolov4-tiny|yolov3-face]-'
'[{dimension}], where dimension could be a single '
'number (e.g. 288, 416, 608) or WxH (e.g. 416x256)'))
args = parser.parse_args()
# if 80 <= 0:
# raise SystemExit('ERROR: bad category_num (%d)!' % 80)
cfg_file_path = '%s.cfg' % model_file
if not os.path.isfile(cfg_file_path):
raise SystemExit('ERROR: file (%s) not found!' % cfg_file_path)
weights_file_path = '%s.weights' % model_file
if not os.path.isfile(weights_file_path):
raise SystemExit('ERROR: file (%s) not found!' % weights_file_path)
output_file_path = '%s.onnx' % model_file
if not verify_classes(model_file, category_num):
raise SystemExit('ERROR: bad category_num (%d)' % category_num)
# Derive yolo input width/height from model name.
# For example, "yolov4-416x256" -> width=416, height=256
w, h = get_input_wh(model_file)
# These are the only layers DarkNetParser will extract parameters
# from. The three layers of type 'yolo' are not parsed in detail
# because they are included in the post-processing later.
supported_layers = [
'net', 'convolutional', 'maxpool', 'shortcut', 'route', 'upsample',
'yolo'
]
# Create a DarkNetParser object, and the use it to generate an
# OrderedDict with all layer's configs from the cfg file.
print('Parsing DarkNet cfg file...')
parser = DarkNetParser(supported_layers)
layer_configs = parser.parse_cfg_file(cfg_file_path)
# We do not need the parser anymore after we got layer_configs.
del parser
# In above layer_config, there are three outputs that we need to
# know the output shape of (in CHW format).
output_tensor_dims = OrderedDict()
c = (80 + 5) * 3
if 'yolov3' in model_file:
if 'tiny' in model_file:
output_tensor_dims['016_convolutional'] = [c, h // 32, w // 32]
output_tensor_dims['023_convolutional'] = [c, h // 16, w // 16]
elif 'spp' in model_file:
output_tensor_dims['089_convolutional'] = [c, h // 32, w // 32]
output_tensor_dims['101_convolutional'] = [c, h // 16, w // 16]
output_tensor_dims['113_convolutional'] = [c, h // 8, w // 8]
elif 'face' in model_file:
output_tensor_dims['082_convolutional'] = [18, h // 32, w // 32]
output_tensor_dims['094_convolutional'] = [18, h // 16, w // 16]
output_tensor_dims['106_convolutional'] = [18, h // 8, w // 8]
else:
output_tensor_dims['082_convolutional'] = [c, h // 32, w // 32]
output_tensor_dims['094_convolutional'] = [c, h // 16, w // 16]
output_tensor_dims['106_convolutional'] = [c, h // 8, w // 8]
elif 'yolov4' in model_file:
if 'tiny' in model_file:
output_tensor_dims['030_convolutional'] = [c, h // 32, w // 32]
output_tensor_dims['037_convolutional'] = [c, h // 16, w // 16]
else:
output_tensor_dims['139_convolutional'] = [c, h // 8, w // 8]
output_tensor_dims['150_convolutional'] = [c, h // 16, w // 16]
output_tensor_dims['161_convolutional'] = [c, h // 32, w // 32]
else:
raise SystemExit('ERROR: unknown model (%s)!' % model_file)
# Create a GraphBuilderONNX object with the specified output tensor
# dimensions.
print('Building ONNX graph...')
builder = GraphBuilderONNX(model_file, output_tensor_dims)
# Now generate an ONNX graph with weights from the previously parsed
# layer configurations and the weights file.
yolo_model_def = builder.build_onnx_graph(
layer_configs=layer_configs,
weights_file_path=weights_file_path,
verbose=True)
# Once we have the model definition, we do not need the builder anymore.
del builder
# Perform a sanity check on the ONNX model definition.
print('Checking ONNX model...')
onnx.checker.check_model(yolo_model_def)
# Serialize the generated ONNX graph to this file.
print('Saving ONNX file...')
onnx.save(yolo_model_def, output_file_path)
print('Done.')
def build_engine(model_name, category_num, do_int8, dla_core, verbose=False):
"""Build a TensorRT engine from ONNX using the older API."""
net_w, net_h = get_input_wh(model_name)
print('Loading the ONNX file...')
onnx_data = load_onnx(model_name)
if onnx_data is None:
return None
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger()
EXPLICIT_BATCH = [] if trt.__version__[0] < '7' else \
[1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)]
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
if do_int8 and not builder.platform_has_fast_int8:
raise RuntimeError('INT8 not supported on this platform')
if not parser.parse(onnx_data):
print('ERROR: Failed to parse the ONNX file.')
for error in range(parser.num_errors):
print(parser.get_error(error))
return None
network = set_net_batch(network, MAX_BATCH_SIZE)
print('Adding yolo_layer plugins...')
network = add_yolo_plugins(
network, model_name, category_num, TRT_LOGGER)
print('Building an engine. This would take a while...')
print('(Use "--verbose" or "-v" to enable verbose logging.)')
if trt.__version__[0] < '7': # older API: build_cuda_engine()
if dla_core >= 0:
raise RuntimeError('DLA core not supported by old API')
builder.max_batch_size = MAX_BATCH_SIZE
builder.max_workspace_size = 1 << 30
builder.fp16_mode = True # alternative: builder.platform_has_fast_fp16
if do_int8:
from calibrator import YOLOEntropyCalibrator
builder.int8_mode = True
builder.int8_calibrator = YOLOEntropyCalibrator(
'calib_images', (net_h, net_w), 'calib_%s.bin' % model_name)
engine = builder.build_cuda_engine(network)
else: # new API: build_engine() with builder config
builder.max_batch_size = MAX_BATCH_SIZE
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30
config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
config.set_flag(trt.BuilderFlag.FP16)
profile = builder.create_optimization_profile()
profile.set_shape(
'000_net', # input tensor name
(MAX_BATCH_SIZE, 3, net_h, net_w), # min shape
(MAX_BATCH_SIZE, 3, net_h, net_w), # opt shape
(MAX_BATCH_SIZE, 3, net_h, net_w)) # max shape
config.add_optimization_profile(profile)
if do_int8:
from calibrator import YOLOEntropyCalibrator
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = YOLOEntropyCalibrator(
'calib_images', (net_h, net_w),
'calib_%s.bin' % model_name)
config.set_calibration_profile(profile)
if dla_core >= 0:
config.default_device_type = trt.DeviceType.DLA
config.DLA_core = dla_core
config.set_flag(trt.BuilderFlag.STRICT_TYPES)
print('Using DLA core %d.' % dla_core)
engine = builder.build_engine(network, config)
if engine is not None:
print('Completed creating engine.')
return engine