转onnx模型

import sys
sys.path.append('./data')
sys.path.append('./model')

from model import MobileNetV3_large


import torch

def convert_onnx():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_path = 'E:/work/ThermalFireRecognition/FireRecognition-test/weights/best3.pt' 
    model = MobileNetV3_large(num_classes=2)


    model.load_state_dict(torch.load(model_path, map_location=device))

    model.to(device)
    model.eval()
    dummy_input = torch.randn(1, 3, 224, 224).to(device)#输入大小   #data type nchw
    onnx_path = 'E:/work/ThermalFireRecognition/FireRecognition-test/weights/final.onnx'
    torch.onnx.export(model, dummy_input, onnx_path, input_names=['input'], output_names=['output'],opset_version=11)
    print('convert retinaface to onnx finish!!!')

if __name__ == "__main__" :
    convert_onnx()

onnx模型简化

## 首先需要将模型转化为onnx
## 加载虚拟环境，并导入 onnx-simplifier 简化网络模型结构

conda activate evns
pip install onnx-simplifier -i https://mirror.baidu.com/pypi/simple
python -m onnxsim ${INPUT_ONNX_MODEL} ${OUTPUT_ONNX_MODEL}
"Sample: python -m onnxsim ./yolov5.onnx ./yolov5_sim.onnx"