mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-01 03:45:06 +00:00
78 lines
3.4 KiB
Python
78 lines
3.4 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
# coding: utf-8
|
||
|
|
# -------------------------------------------------------------------------
|
||
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||
|
|
# Licensed under the MIT License. See License.txt in the project root for
|
||
|
|
# license information.
|
||
|
|
# --------------------------------------------------------------------------
|
||
|
|
|
||
|
|
import unittest
|
||
|
|
import onnx
|
||
|
|
import onnxruntime
|
||
|
|
import numpy as np
|
||
|
|
from onnx import helper, TensorProto, numpy_helper
|
||
|
|
from onnxruntime.quantization.onnx_model import ONNXModel
|
||
|
|
from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count, check_op_type_order
|
||
|
|
|
||
|
|
|
||
|
|
def generate_input_initializer(tensor_shape, tensor_dtype, input_name):
|
||
|
|
'''
|
||
|
|
Helper function to generate initializers for test inputs
|
||
|
|
'''
|
||
|
|
tensor = np.random.normal(0, 0.3, tensor_shape).astype(tensor_dtype)
|
||
|
|
init = numpy_helper.from_array(tensor, input_name)
|
||
|
|
return init
|
||
|
|
|
||
|
|
class TestONNXModel(unittest.TestCase):
|
||
|
|
def construct_model(self, model_path):
|
||
|
|
# (input)
|
||
|
|
# |
|
||
|
|
# GRU
|
||
|
|
# / \
|
||
|
|
# Conv(1) \
|
||
|
|
# | \
|
||
|
|
# Relu Conv(2)
|
||
|
|
# | |
|
||
|
|
# \ /
|
||
|
|
# Add
|
||
|
|
# |
|
||
|
|
# (output)
|
||
|
|
initializers = []
|
||
|
|
input = helper.make_tensor_value_info('input', TensorProto.FLOAT, [4, 8, 12])
|
||
|
|
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [4, 2, 8, 8])
|
||
|
|
|
||
|
|
# make GRU
|
||
|
|
initializers.append(generate_input_initializer([2, 24, 12], np.float32, 'W_GRU'))
|
||
|
|
initializers.append(generate_input_initializer([2, 24, 8], np.float32, 'R_GRU'))
|
||
|
|
initializers.append(generate_input_initializer([2, 8, 8], np.float32, 'H_GRU'))
|
||
|
|
gru_node = onnx.helper.make_node(
|
||
|
|
'GRU',
|
||
|
|
['input', 'W_GRU', 'R_GRU', '', '', 'H_GRU'],
|
||
|
|
['GRU_O'],
|
||
|
|
hidden_size = 8,
|
||
|
|
direction = 'bidirectional')
|
||
|
|
|
||
|
|
initializers.append(generate_input_initializer([2, 2, 1, 1], np.float32, 'W1'))
|
||
|
|
initializers.append(generate_input_initializer([2, 2, 1, 1], np.float32, 'W2'))
|
||
|
|
initializers.append(generate_input_initializer([2], np.float32, 'B1'))
|
||
|
|
initializers.append(generate_input_initializer([2], np.float32, 'B2'))
|
||
|
|
conv_node_1 = onnx.helper.make_node('Conv', ['GRU_O', 'W1', 'B1'], ['Conv1_O'], name='Conv1')
|
||
|
|
conv_node_2 = onnx.helper.make_node('Conv', ['GRU_O', 'W2', 'B2'], ['Conv2_O'], name='Conv2')
|
||
|
|
relu_node = onnx.helper.make_node('Relu', ['Conv1_O'], ['Relu_O'], name='Relu')
|
||
|
|
add_node = onnx.helper.make_node('Add', ['Relu_O', 'Conv2_O'], ['output'], name='Add')
|
||
|
|
graph = helper.make_graph([conv_node_1, relu_node, conv_node_2, gru_node, add_node],
|
||
|
|
'onnx_model_test', [input], [output], initializer=initializers)
|
||
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
||
|
|
onnx.save(model, model_path)
|
||
|
|
|
||
|
|
def test_topo_sort(self):
|
||
|
|
test_model_path = 'onnx_model_topo_sort.onnx'
|
||
|
|
self.construct_model(test_model_path)
|
||
|
|
onnx_model = ONNXModel(onnx.load(test_model_path))
|
||
|
|
check_op_type_order(self, onnx_model.model, ['Conv', 'Relu', 'Conv', 'GRU', 'Add'])
|
||
|
|
onnx_model.topological_sort()
|
||
|
|
check_op_type_order(self, onnx_model.model, ['GRU', 'Conv', 'Conv', 'Relu', 'Add'])
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|