onnxruntime 导出网络每层数据

最新推荐文章于 2023-07-01 11:26:43 发布

Aisinjulo

最新推荐文章于 2023-07-01 11:26:43 发布

阅读量331

点赞数

文章标签：深度学习 pytorch

本文链接：https://blog.csdn.net/Aisinjulo/article/details/128089616

版权

import os
import onnx
from PIL import Image
from torchvision import transforms
import copy
import numpy as np
import logging
import onnxruntime
from collections import OrderedDict
from onnx import shape_inference
logging.basicConfig(level=logging.INFO)
from onnx import shape_inference, TensorProto, version_converter, numpy_helper
logger = logging.getLogger("[ONNXOPTIMIZER]")

def test_model_by_onnxruntime(model):
    logger.info("Test model by onnxruntime")

    # genarate random number
    input_shape = model.graph.input[0].type.tensor_type.shape.dim
    image_shape = [x.dim_value for x in input_shape]
    image_shape_new = []
    for x in image_shape:
        if x == 0:
            image_shape_new.append(1)
        else:
            image_shape_new.append(x)
    image_shape = image_shape_new
    img_array = np.array(np.random.random(image_shape), dtype = np.float32)
    img = img_array

    # load image
    # input_image = Image.open("./images/ILSVRC2012_val_00000001_n01751748.JPEG")
    # preprocess = transforms.Compose([
    #     transforms.Resize(256),
    #     transforms.CenterCrop(224),
    #     transforms.ToTensor(),
    #     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    # ])
    # input_tensor = preprocess(input_image)
    # input_batch = input_tensor.unsqueeze(0)
    # img = input_batch.numpy()

    for node in model.graph.node:
        for output in node.output:
            model.graph.output.extend([onnx.ValueInfoProto(name=output)])
    ort_session = onnxruntime.InferenceSession(model.SerializeToString(),
                                     providers=['CUDAExecutionProvider'])
    ort_inputs = {}
    for i, input_ele in enumerate(ort_session.get_inputs()):
        ort_inputs[input_ele.name] = img

    outputs = [x.name for x in ort_session.get_outputs()]
    ort_outs = ort_session.run(outputs, ort_inputs)
    # ort_outs = OrderedDict(zip(outputs, ort_outs))

    for i in range(len(ort_outs)):
        ort_outs[i] = ort_outs[i].astype(np.float16)
        ort_outs[i].tofile("./layers_result/layer_"+str(i)+".bin")   

    logger.info("Test model by onnxruntime success")
    return ort_outs

onnx_model = onnx.load("mobilenet_v2.onnx")
ort_outs = test_model_by_onnxruntime(onnx_model)