diff options
Diffstat (limited to 'model-integration/src/test/models/onnx/simple/simple.py')
-rwxr-xr-x | model-integration/src/test/models/onnx/simple/simple.py | 32 |
1 files changed, 32 insertions, 0 deletions
diff --git a/model-integration/src/test/models/onnx/simple/simple.py b/model-integration/src/test/models/onnx/simple/simple.py new file mode 100755 index 00000000000..4471ed812b8 --- /dev/null +++ b/model-integration/src/test/models/onnx/simple/simple.py @@ -0,0 +1,32 @@ +import onnx +from onnx import helper, TensorProto + +QUERY_TENSOR = helper.make_tensor_value_info('query_tensor', TensorProto.FLOAT, [1, 4]) +ATTRIBUTE_TENSOR = helper.make_tensor_value_info('attribute_tensor', TensorProto.FLOAT, [4, 1]) +BIAS_TENSOR = helper.make_tensor_value_info('bias_tensor', TensorProto.FLOAT, [1, 1]) +OUTPUT = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 1]) + +nodes = [ + helper.make_node( + 'MatMul', + ['query_tensor', 'attribute_tensor'], + ['matmul'], + ), + helper.make_node( + 'add', + ['matmul', 'bias_tensor'], + ['output'], + ), +] +graph_def = helper.make_graph( + nodes, + 'simple_scoring', + [ + QUERY_TENSOR, + ATTRIBUTE_TENSOR, + BIAS_TENSOR, + ], + [OUTPUT], +) +model_def = helper.make_model(graph_def, producer_name='simple.py') +onnx.save(model_def, 'simple.onnx') |