Add onnx2ncnn tool (#276)

* Add onnx2ncnn tool

* fix docs

* fix lint
This commit is contained in:
q.yao 2021-12-15 18:43:09 +08:00 committed by GitHub
parent b1a202ef04
commit 43256977fa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 66 additions and 0 deletions

View File

@ -84,3 +84,24 @@ python tools/onnx2tensorrt.py \
- `--device-id` : The device index, default to `0`.
- `--calib-file` : The calibration data used to calibrate engine to int8.
- `--log-level` : To set log level which in `'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'`. If not specified, it will be set to `INFO`.
## onnx2ncnn
This tool helps to convert an `ONNX` model to an `ncnn` model.
### Usage
```bash
python tools/onnx2ncnn.py \
${ONNX_PATH} \
${NCNN_PARAM} \
${NCNN_BIN} \
--log-level INFO
```
### Description of all arguments
- `onnx_path` : The path of the `ONNX` model to convert from.
- `output_param` : The converted `ncnn` param path.
- `output_bin` : The converted `ncnn` bin path.
- `--log-level` : To set log level which in `'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'`. If not specified, it will be set to `INFO`.

45
tools/onnx2ncnn.py Normal file
View File

@ -0,0 +1,45 @@
import argparse
import logging
from mmdeploy.apis.ncnn import onnx2ncnn
def parse_args():
parser = argparse.ArgumentParser(description='Convert ONNX to ncnn.')
parser.add_argument('onnx_path', help='ONNX model path')
parser.add_argument('output_param', help='output ncnn param path')
parser.add_argument('output_bin', help='output bin path')
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()
logging.basicConfig(
format='%(asctime)s,%(name)s %(levelname)-8s'
' [%(filename)s:%(lineno)d] %(message)s',
datefmt='%Y-%m-%d:%H:%M:%S')
logger = logging.getLogger()
logger.setLevel(args.log_level)
onnx_path = args.onnx_path
output_param = args.output_param
output_bin = args.output_bin
logging.info(f'onnx2ncnn: \n\tonnx_path: {onnx_path} ')
try:
onnx2ncnn(onnx_path, output_param, output_bin)
logging.info('onnx2ncnn success.')
except Exception as e:
logging.error(e)
logging.error('onnx2ncnn failed.')
if __name__ == '__main__':
main()