diff options
author | Håvard Pettersen <havardpe@oath.com> | 2022-02-10 11:44:36 +0000 |
---|---|---|
committer | Håvard Pettersen <havardpe@oath.com> | 2022-02-10 11:44:36 +0000 |
commit | f4fd8f8e29fa87ec79afce4172ce3d72ab6693f0 (patch) | |
tree | 02fd48553341f4c6da973e0cd7e5a818e4461cd8 /eval/src/tests/tensor/onnx_wrapper/probe_model.py | |
parent | fed60ae90b4e230129b4e9e6b72dd125d354e6c0 (diff) |
add code to simplify onnx model testing
test model that combines probing and inference of dimension sizes
Diffstat (limited to 'eval/src/tests/tensor/onnx_wrapper/probe_model.py')
-rwxr-xr-x | eval/src/tests/tensor/onnx_wrapper/probe_model.py | 35 |
1 files changed, 35 insertions, 0 deletions
diff --git a/eval/src/tests/tensor/onnx_wrapper/probe_model.py b/eval/src/tests/tensor/onnx_wrapper/probe_model.py new file mode 100755 index 00000000000..529fa23b2b1 --- /dev/null +++ b/eval/src/tests/tensor/onnx_wrapper/probe_model.py @@ -0,0 +1,35 @@ +# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +import onnx +from onnx import helper, TensorProto + +IN1 = helper.make_tensor_value_info('in1', TensorProto.FLOAT, [-1, 'inner']) +IN2 = helper.make_tensor_value_info('in2', TensorProto.FLOAT, ['outer', -1]) +OUT1 = helper.make_tensor_value_info('out1', TensorProto.FLOAT, [-1, 'inner']) +OUT2 = helper.make_tensor_value_info('out2', TensorProto.FLOAT, ['outer', -1]) +OUT3 = helper.make_tensor_value_info('out3', TensorProto.FLOAT, [-1, -1]) + +nodes = [ + helper.make_node( + 'Add', + ['in1', 'in2'], + ['out1'], + ), + helper.make_node( + 'Sub', + ['in1', 'in2'], + ['out2'], + ), + helper.make_node( + 'Mul', + ['in1', 'in2'], + ['out3'], + ), +] +graph_def = helper.make_graph( + nodes, + 'probe_model', + [IN1, IN2], + [OUT1, OUT2, OUT3], +) +model_def = helper.make_model(graph_def, producer_name='probe_model.py', opset_imports=[onnx.OperatorSetIdProto(version=12)]) +onnx.save(model_def, 'probe_model.onnx') |