onnx操作 云朵喵喵

1.pt文件转换onnx

import torch
from models.experimental import attempt_load

#加载模型权重
model = attempt_load("runs/train/exp/weights/best.pt",device=torch.device("cpu"))
#设置为评估模式
model.eval()
#准备一个示例输入
input_tensor = torch.randn(1,3,640,640)
#导出模型
torch.onnx.export(model,input_tensor,"yolov5s.onnx")

2.加载onnx

import onnx


# 加载模型
model = onnx.load("quantized_best.onnx")
# 检测模型结构
onnx.checker.check_model(model)

3.onnx量化

# 加载 ONNX 模型
from onnxruntime.quantization import quantize_dynamic, QuantType

onnx_path = r'runs/train/exp2/weights/best.onnx'
# 动态量化
quantized_model = quantize_dynamic(onnx_path, 'quantized_best.onnx', weight_type=QuantType.QUInt8)