101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
import argparse
|
|
import os
|
|
import subprocess
|
|
|
|
# list of tuple: config, pretrained model, onnx filename
|
|
CONFIGS = [
|
|
(
|
|
'mmclassification/configs/vision_transformer/vit-base-p32_ft-64xb64_in1k-384.py', # noqa: E501
|
|
'https://download.openmmlab.com/mmclassification/v0/vit/finetune/vit-base-p32_in21k-pre-3rdparty_ft-64xb64_in1k-384_20210928-9cea8599.pth', # noqa: E501
|
|
'vit.onnx'),
|
|
(
|
|
'mmclassification/configs/resnet/resnet50_8xb32_in1k.py',
|
|
'https://download.openmmlab.com/mmclassification/v0/resnet/resnet50_8xb32_in1k_20210831-ea4938fc.pth', # noqa: E501
|
|
'resnet50.onnx',
|
|
),
|
|
(
|
|
'mmclassification/configs/resnet/resnet18_8xb32_in1k.py',
|
|
'https://download.openmmlab.com/mmclassification/v0/resnet/resnet18_8xb32_in1k_20210831-fbbb1da6.pth', # noqa: E501
|
|
'resnet18.onnx',
|
|
'https://media.githubusercontent.com/media/tpoisonooo/mmdeploy-onnx2ncnn-testdata/main/resnet18.onnx', # noqa: E501
|
|
),
|
|
(
|
|
'mmclassification/configs/mobilenet_v2/mobilenet-v2_8xb32_in1k.py',
|
|
'https://download.openmmlab.com/mmclassification/v0/mobilenet_v2/mobilenet_v2_batch256_imagenet_20200708-3b2dc3af.pth', # noqa: E501
|
|
'mobilenet-v2.onnx',
|
|
'https://media.githubusercontent.com/media/tpoisonooo/mmdeploy-onnx2ncnn-testdata/main/mobilenet-v2.onnx', # noqa: E501
|
|
)
|
|
]
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
description='MMDeploy onnx2ncnn test tool.')
|
|
parser.add_argument(
|
|
'--run', type=bool, help='Execute mmdeploy_onnx2ncnn bin.')
|
|
parser.add_argument(
|
|
'--repo-dir', type=str, default='~/', help='mmcls directory.')
|
|
parser.add_argument(
|
|
'--out',
|
|
type=str,
|
|
default='onnx_output',
|
|
help='onnx model output directory.')
|
|
parser.add_argument(
|
|
'--generate-onnx', type=bool, help='Generate onnx model.')
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def generate_onnx(args):
|
|
import mmengine
|
|
mmengine.mkdir_or_exist(args.out)
|
|
for conf in CONFIGS:
|
|
config = os.path.join(args.repo_dir, conf[0])
|
|
model = conf[1]
|
|
convert_cmd = [
|
|
'python', 'tools/deploy.py',
|
|
'configs/mmcls/classification_ncnn_static.py', config, model,
|
|
'cat-dog.png', '--work-dir', 'work_dir', '--device', 'cpu'
|
|
]
|
|
print(subprocess.call(convert_cmd))
|
|
|
|
move_cmd = [
|
|
'mv', 'work_dir/end2end.onnx',
|
|
os.path.join(args.out, conf[2])
|
|
]
|
|
print(subprocess.call(move_cmd))
|
|
|
|
|
|
def run(args):
|
|
for conf in CONFIGS:
|
|
if len(conf) < 4:
|
|
continue
|
|
download_url = conf[3]
|
|
filename = conf[2]
|
|
download_cmd = ['wget', download_url]
|
|
# show processbar
|
|
os.system(' '.join(download_cmd))
|
|
|
|
convert_cmd = [
|
|
'./mmdeploy_onnx2ncnn', filename, 'onnx.param', 'onnx.bin'
|
|
]
|
|
subprocess.run(convert_cmd, capture_output=True, check=True)
|
|
|
|
|
|
def main():
|
|
"""test `onnx2ncnn.cpp`
|
|
|
|
First generate onnx model then convert it with `mmdeploy_onnx2ncnn`.
|
|
"""
|
|
args = parse_args()
|
|
if args.generate_onnx:
|
|
generate_onnx(args)
|
|
if args.run:
|
|
run(args)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|