pytorch框架--网络方面--onnx模型-转换-保存-加载-部署

模型.pth 转换为 .onnx

# 首先正常加载模型model
# 根据图像需求输入
dummy_input = torch.randn(1, 3, 224, 224)
# 保存.onnx格式模型  
torch.onnx._export(model, dummy_input, "cnn.onnx", verbose=True)

测试onnx模型是否可用

# 测试onnx模型，需安装onnx
pip install onnx

测试代码：

import onnx
# 验证模型onnx格式
model = onnx.load('cnn.onnx')
onnx.checker.check_model(model)

加载测试模型

# 加模型必备（CPU） 如果使用GPU，目前还没尝试
pip install onnxruntime

加载代码：

import onnxruntime
# 加载模型
ort_session = onnxruntime.InferenceSession("cnn.onnx")
# 输入图像 注意：图像需要预处理，还需转np.array格式
ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(img1)

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

ort_inputs = {ort_session.get_inputs()[0].name: ortvalue}
ort_outs = ort_session.run(None, ort_inputs)


dist = {0: "猫", 1: "狗"}
# 转numpy格式,列表内取第一个
a1 = dist[ort_outs[0][0].argmax()]
print(a1)