Skip to content

Commit

Permalink
Add onnx2ncnn tool (open-mmlab#276)
Browse files Browse the repository at this point in the history
* Add onnx2ncnn tool

* fix docs

* fix lint
  • Loading branch information
grimoire authored Dec 15, 2021
1 parent b1a202e commit 4325697
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 0 deletions.
21 changes: 21 additions & 0 deletions docs/useful_tools.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 changes: 45 additions & 0 deletions tools/onnx2ncnn.py
Original file line number Diff line number Diff line change
@@ -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()

0 comments on commit 4325697

Please sign in to comment.