onnx篇---验证onnx转换成功与否

1. 调用opencv 来查看网路输出的结果

use_opencv.py

import cv2
import onnxruntime
import numpy as np


onnx_path_demo = '../data/weights/Yolo_bin.onnx'
# onnx_path_demo = 'YOLO.onnx'
image_path = '../data/img/0601_Bin_061.jpg'

image_src = cv2.imread(image_path)
resized = cv2.resize(image_src, (416, 416), interpolation=cv2.INTER_LINEAR)
img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
img_in = np.expand_dims(img_in, axis=0)
img_in /= 255.0

session = onnxruntime.InferenceSession(onnx_path_demo)
input_name = session.get_inputs()[0].name
outputs = session.run(None, {input_name: img_in})

print(outputs[0].shape, outputs[1].shape, outputs[2].shape)

得到网络的输出结果，验证权重转换的准确性

2. 使用onnx加载.onnx权重

validation_onnx.py
查看输出结果以及网路输出shape

# -- coding: utf-8 --
import torch
import numpy as np
from PIL import Image
import onnxruntime
from onnxruntime.datasets import get_example
from model.nets.yolo4 import YoloBody


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

# 测试数据
dummy_input = torch.randn(1, 3, 416, 416, device='cpu')

example_model = get_example(r'D:\second_project\project\yolo\object-detections-yolov4-bin-3\data\weights\Yolo_bin.onnx')
# netron.start(example_model) 使用 netron python 包可视化网络
sess = onnxruntime.InferenceSession(example_model)

# onnx 网络输出
onnx_out = sess.run(None, {'input': to_numpy(dummy_input)})
print(onnx_out)
print(onnx_out[0].shape)
print(onnx_out[1].shape)
print(onnx_out[2].shape)

3. 加载原始网络权重

查看输出结果以及网络输出shape

# ---------------------------------------------------#
#   对比验证结果
# ---------------------------------------------------#
weight_file = '../data/weights/Epoch75-Total_Loss0.7506-Val_Loss1.6806.pth'
num_anchors, num_classes = 3, 3
model = YoloBody(num_anchors, num_classes)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 加快模型训练的效率
model_dict = model.state_dict()
pretrained_dict = torch.load(weight_file, map_location=device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

model.eval()
with torch.no_grad():
    # pytorch model 网络输出
    torch_out = model(dummy_input)
    print(torch_out)
    print(torch_out[0].shape)
    print(torch_out[1].shape)
    print(torch_out[2].shape)