onnx推理保存每一层输出保存在本地

背景

在模型转换的时候，发现转换后的模型精度与训练时候不一致，为了弄明白到底那一层精度丢失比较大，于是把每一层的输出保留下来，进行对比。

onnx推理代码并保存结果

import os
import cv2
import onnx
import onnxruntime
import numpy as np
from collections import OrderedDict

def get_img_data(img_path):
    
    img = cv2.imread(img_path)
    img = cv2.resize(img, (864, 480))
    img = np.transpose(img, (2, 0, 1)).astype(np.float32)
    img = np.expand_dims(img, axis = 0)
    img = img / 255.0
    
    return img


def get_infere_data(onnx_path, img):
    
    onnx_model = onnx.load(onnx_path)
    for node in onnx_model.graph.node:
        for out_node in node.output:
            onnx_model.graph.output.extend([onnx.ValueInfoProto(name=out_node)])
            
    ort_seesion = onnxruntime.InferenceSession(onnx_model.SerializeToString())
    
    ort_input = {ort_seesion.get_inputs()[0].name: img}
    
    
    outputs = [x.name for x in ort_seesion.get_outputs()]
    
    ort_outs = ort_seesion.run(outputs, ort_input)
    ort_outs = OrderedDict(zip(outputs, ort_outs))
    
    return ort_outs

def write_output_txt(output, txt_file):
    
    names = output.keys()
    for name in names:
        txt_path = os.path.join(txt_file, name + '.txt')
        with open(txt_path, 'a') as f:
            new_data = output[name].reshape(-1)
            for i in range(new_data.shape[0]):
                line = str(new_data[i])
                line += '\n'
                f.write(line)

if __name__ == "__main__":
    onnx_path = '～/weights/best-sim.onnx'
    img_path  = '～/img/test.bmp'
    txt_file  = '～/out_txt/'
    img = get_img_data(img_path)
    output = get_infere_data(onnx_path, img)
    write_output_txt(output, txt_file)

总结

根据保存的结果与转换之前的结果做对比，就知道那一层差异比较大了。