# 首先正常加载模型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)
