onnx mxnet输出层结果对比

最新推荐文章于 2024-05-15 16:00:33 发布

zhaosuyuan

最新推荐文章于 2024-05-15 16:00:33 发布

阅读量76

点赞数

分类专栏： onnx转换 mxnet 文章标签： mxnet python 人工智能

本文链接：https://blog.csdn.net/zhaosuyuan/article/details/131416875

版权

onnx转换同时被 2 个专栏收录

4 篇文章 0 订阅

订阅专栏

mxnet

2 篇文章 0 订阅

订阅专栏

onnx和mxnet转换输出层结果及其cos相似度计算对比


import cv2
import os, sys
import argparse
import numpy as np
import onnxruntime
import mxnet as mx
from collections import namedtuple

class onnxPredictor:
    def __init__(self, path):
        self.session = onnxruntime.InferenceSession(path, None)
        self.input_shape = tuple(self.session.get_inputs()[0].shape)
        self.outputs = list(map(lambda x:x.name, self.session.get_outputs()))

    def forward(self, tensor):
        out = self.session.run(self.outputs, {"data": tensor})
        return out

Batch = namedtuple('Batch',['data'])
class mxnetPredictor:
    def __init__(self, prefix, epoch, input_shape, ctx=mx.cpu()):
        sym, arg_params, aux_params = mx.model.load_checkpoint(prefix=prefix, epoch=epoch)
        self.input_shape = input_shape
        self.outputs = sym.list_outputs()
        self.mod = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
        self.mod.bind(for_training=False, data_shapes=[('data', self.input_shape)])
        self.mod.set_params(arg_params, aux_params, allow_missing=False, allow_extra=False)

    def forward(self, tensor):
        self.mod.forward(Batch([mx.nd.array(tensor)]))
        out = self.mod.get_outputs()
        out = list(map(lambda x:x.asnumpy(), out))
        return out

def compare_cos_sim(tensor1, tensor2):
    t1 = tensor1.reshape(-1)
    t2 = tensor2.reshape(-1)
    n1 = np.linalg.norm(t1)
    n2 = np.linalg.norm(t2)
    cos = np.dot(t1/n1, t2.T/n2)
    print('cos',cos)
    '''
    if cos < 0.8 or True:
        print("\t", t1[:10])
        print("\t", t2[:10])
    print("\t", n1, n2, cos)
    '''    
    
def img2tensor(mat):
    if mat is None:
        tensor = np.random.randint(0, 256, input_shape, dtype=np.uint8)
    else:
        mat = cv2.resize(mat, (input_shape[3], input_shape[2]))
        if input_shape[1] == 1:
            mat = cv2.cvtColor(mat, cv2.COLOR_BGR2GRAY)
            mat = mat[:,:,np.newaxis]
        tensor = np.expand_dims(np.transpose(mat, (2,0,1)), axis=0)
    tensor = tensor.astype(np.float32)
    return tensor
    
def compare_models(path_onnx, path_mxnet, mat):
    netonnx = onnxPredictor(path_onnx)
    input_shape = netonnx.input_shape
    netmxnet = mxnetPredictor(path_mxnet, 0, input_shape)
    tensor = img2tensor(mat)

    out_onnx = netonnx.forward(tensor)
    out_mxnet = netmxnet.forward(tensor)
    #print('out_onnx***************',out_onnx)
    #print('out_mxnet================',out_mxnet)

    for out in out_onnx:
        print('onnx max',np.max(out),out.shape)
    print('====================================')        
    for out in out_mxnet:
        print('mxnet max',np.max(out),out.shape)

    outputs = netonnx.outputs
    #print(outputs)
    for i, output in enumerate(outputs):
        compare_cos_sim(out_onnx[i], out_mxnet[i])

def parse_args():
  parser = argparse.ArgumentParser(description="check results")
  parser.add_argument("--onnx", required=True, help="path to onnx model")
  parser.add_argument("--mxnet", required=True, help="path to mxnet model")
  parser.add_argument("--input", type=str, default="")
  args = parser.parse_args()
  return args

if __name__ == "__main__":
    args = parse_args()
    path_onnx = args.onnx
    path_mxnet = args.mxnet
    mat = None
    if len(args.input) != 0:
        print("imread:", args.input)
        mat = cv2.imread(args.input)
    compare_models(path_onnx, path_mxnet, mat)