diff --git a/example/onnx-runtime/FastestDet.onnx b/example/onnx-runtime/FastestDet.onnx index 0666a12..cce8577 100644 Binary files a/example/onnx-runtime/FastestDet.onnx and b/example/onnx-runtime/FastestDet.onnx differ diff --git a/requirements.txt b/requirements.txt index 1f7df40..c4ddab8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,8 @@ -numpy==1.19.5 +numpy==1.23.0 +onnx==1.12.0 +onnx_simplifier==0.3.10 +onnxruntime==1.11.1 +onnxsim==0.4.0 opencv_python==4.2.0.34 pycocotools==2.0.4 PyYAML==6.0 diff --git a/test.py b/test.py index c0d9fb7..53ca3dd 100644 --- a/test.py +++ b/test.py @@ -1,7 +1,9 @@ import os import cv2 +import onnx import time import argparse +from onnxsim import simplify import torch from utils.tool import * @@ -55,12 +57,18 @@ if __name__ == '__main__': # 导出onnx模型 if opt.onnx: - torch.onnx.export(model, #model being run - img, # model input (or a tuple for multiple inputs) - "./FastestDet.onnx", # where to save the model (can be a file or file-like object) + torch.onnx.export(model, # model being run + img, # model input (or a tuple for multiple inputs) + "./FastestDet.onnx", # where to save the model (can be a file or file-like object) export_params=True, # store the trained parameter weights inside the model file opset_version=11, # the ONNX version to export the model to do_constant_folding=True) # whether to execute constant folding for optimization + # onnx-sim + onnx_model = onnx.load("./FastestDet.onnx") # load onnx model + model_simp, check = simplify(onnx_model) + assert check, "Simplified ONNX model could not be validated" + print("onnx sim sucess...") + onnx.save(model_simp, "./FastestDet.onnx") # 导出torchscript模型 if opt.torchscript: