60 lines
1.8 KiB
Python
60 lines
1.8 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
|
|
import os
|
|
import subprocess
|
|
|
|
# list of tuple: onnx filename, download_url and model_config.
|
|
|
|
CONFIGS = [
|
|
[
|
|
'hrnet.onnx',
|
|
'https://media.githubusercontent.com/media/tpoisonooo/mmdeploy-onnx2ncnn-testdata/main/hrnet.onnx', # noqa: E501
|
|
'~/mmclassification/configs/hrnet/hrnet-w18_4xb32_in1k.py',
|
|
],
|
|
[
|
|
'resnet18.onnx',
|
|
'https://media.githubusercontent.com/media/tpoisonooo/mmdeploy-onnx2ncnn-testdata/main/resnet18.onnx', # noqa: E501
|
|
'~/mmclassification/configs/resnet/resnet18_8xb16_cifar10.py',
|
|
],
|
|
[
|
|
'mobilenet-v2.onnx',
|
|
'https://media.githubusercontent.com/media/tpoisonooo/mmdeploy-onnx2ncnn-testdata/main/mobilenet-v2.onnx', # noqa: E501
|
|
'~/mmclassification/configs/mobilenet_v2/mobilenet-v2_8xb32_in1k.py',
|
|
]
|
|
]
|
|
|
|
|
|
def prepare_dataset():
|
|
DATASET = (
|
|
'dataset',
|
|
'https://media.githubusercontent.com/media/tpoisonooo/mmdeploy-onnx2ncnn-testdata/main/dataset.tar' # noqa: E501
|
|
) # noqa: E501
|
|
os.system('wget {}'.format(DATASET[1]))
|
|
os.system('tar xvf dataset.tar')
|
|
return DATASET[0]
|
|
|
|
|
|
def main():
|
|
"""test `tools/onnx2ncnn_quant_table.py`
|
|
|
|
First quantize onnx model to ncnn with ppq.
|
|
"""
|
|
data_dir = prepare_dataset()
|
|
|
|
for conf in CONFIGS:
|
|
model = conf[0]
|
|
os.system('wget {}'.format(conf[1]))
|
|
model_cfg = conf[2]
|
|
deploy_cfg = 'configs/mmcls/classification_ncnn-int8_static.py'
|
|
quant_cmd = [
|
|
'python3', 'tools/onnx2ncnn_quant_table.py', '--onnx', model,
|
|
'--deploy-cfg', deploy_cfg, '--model-cfg', model_cfg, '--out-onnx',
|
|
'quant.onnx', '--out-table', 'ncnn.table', '--image-dir', data_dir
|
|
]
|
|
print(' '.join(quant_cmd))
|
|
print(subprocess.call(quant_cmd))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|