mirror of
https://github.com/open-mmlab/mmdeploy.git
synced 2025-01-14 08:09:43 +08:00
b32fc41bed
* [refactor][API2.0] Add onnx export and jit trace (#419) * first commit * add async call * add new api onnx export and jit trace * add decorator * fix ci * fix torchscript ci * fix loader * better pipemanager * remove comment, better import * add kwargs * remove comment * better pipeline manager * remove print * [Refactor][API2.0] Api partition calibration (#433) * first commit * add async call * add new api onnx export and jit trace * add decorator * fix ci * fix torchscript ci * fix loader * better pipemanager * remove comment, better import * add partition * move calibration * Better create_calib_table * better deploy * add kwargs * remove comment * better pipeline manager * rename api, remove reduant variable, and misc * [Refactor][API2.0] Api ncnn openvino (#435) * first commit * add async call * add new api onnx export and jit trace * add decorator * fix ci * fix torchscript ci * fix loader * better pipemanager * remove comment, better import * add ncnn api * finish ncnn api * add openvino support * add kwargs * remove comment * better pipeline manager * merge fix * merge util and onnx2ncnn * fix docstring * [Refactor][API2.0] API for TensorRT (#519) * first commit * add async call * add new api onnx export and jit trace * add decorator * fix ci * fix torchscript ci * fix loader * better pipemanager * remove comment, better import * add partition * move calibration * Better create_calib_table * better deploy * add kwargs * remove comment * Add tensorrt API * better pipeline manager * add tensorrt new api * remove print * rename api, remove reduant variable, and misc * add docstring * [Refactor][API2.0] Api ppl other (#528) * first commit * add async call * add new api onnx export and jit trace * add decorator * fix ci * fix torchscript ci * fix loader * better pipemanager * remove comment, better import * add kwargs * Add new APIS for pplnn sdk and misc * remove comment * better pipeline manager * merge fix * update tools/onnx2pplnn.py * rename function
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import argparse
|
|
import logging
|
|
|
|
from mmdeploy.backend.tensorrt import from_onnx
|
|
from mmdeploy.backend.tensorrt.utils import get_trt_log_level
|
|
from mmdeploy.utils import (get_common_config, get_model_inputs,
|
|
get_root_logger, load_config)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description='Convert ONNX to TensorRT.')
|
|
parser.add_argument('deploy_cfg', help='deploy config path')
|
|
parser.add_argument('onnx_path', help='ONNX model path')
|
|
parser.add_argument('output_prefix', help='output TensorRT engine prefix')
|
|
parser.add_argument('--device-id', help='`the CUDA device id', default=0)
|
|
parser.add_argument(
|
|
'--calib-file',
|
|
help='`the calibration data used to calibrate engine to int8',
|
|
default=None)
|
|
parser.add_argument(
|
|
'--log-level',
|
|
help='set log level',
|
|
default='INFO',
|
|
choices=list(logging._nameToLevel.keys()))
|
|
args = parser.parse_args()
|
|
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
logger = get_root_logger(log_level=args.log_level)
|
|
|
|
deploy_cfg_path = args.deploy_cfg
|
|
deploy_cfg = load_config(deploy_cfg_path)[0]
|
|
onnx_path = args.onnx_path
|
|
output_prefix = args.output_prefix
|
|
device_id = args.device_id
|
|
calib_file = args.calib_file
|
|
|
|
model_id = 0
|
|
common_params = get_common_config(deploy_cfg)
|
|
model_params = get_model_inputs(deploy_cfg)[model_id]
|
|
|
|
final_params = common_params
|
|
final_params.update(model_params)
|
|
|
|
int8_param = final_params.get('int8_param', dict())
|
|
|
|
if calib_file is not None:
|
|
int8_param['calib_file'] = calib_file
|
|
# do not support partition model calibration for now
|
|
int8_param['model_type'] = 'end2end'
|
|
|
|
logger.info(f'onnx2tensorrt: \n\tonnx_path: {onnx_path} '
|
|
f'\n\tdeploy_cfg: {deploy_cfg_path}')
|
|
try:
|
|
from_onnx(
|
|
onnx_path,
|
|
output_prefix,
|
|
input_shapes=final_params['input_shapes'],
|
|
log_level=get_trt_log_level(),
|
|
fp16_mode=final_params.get('fp16_mode', False),
|
|
int8_mode=final_params.get('int8_mode', False),
|
|
int8_param=int8_param,
|
|
max_workspace_size=final_params.get('max_workspace_size', 0),
|
|
device_id=device_id)
|
|
|
|
logger.info('onnx2tensorrt success.')
|
|
except Exception as e:
|
|
logger.error(e)
|
|
logger.error('onnx2tensorrt failed.')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|