diff options
author | Håvard Pettersen <havardpe@oath.com> | 2020-09-03 14:31:06 +0000 |
---|---|---|
committer | Håvard Pettersen <havardpe@oath.com> | 2020-09-05 14:54:13 +0000 |
commit | b5b4b47229990035260806ed644c0f46d0c8a282 (patch) | |
tree | f589e6fc4b288677c69ccbbe000e59f2a6d820ee /eval/src/tests/tensor/onnx_wrapper/int_types.py | |
parent | ec0e8716abbf29d958e768d40b57b23002221c4a (diff) |
adapt and convert between vespa and onnx types
Diffstat (limited to 'eval/src/tests/tensor/onnx_wrapper/int_types.py')
-rwxr-xr-x | eval/src/tests/tensor/onnx_wrapper/int_types.py | 33 |
1 files changed, 33 insertions, 0 deletions
diff --git a/eval/src/tests/tensor/onnx_wrapper/int_types.py b/eval/src/tests/tensor/onnx_wrapper/int_types.py new file mode 100755 index 00000000000..cd82bfd44b5 --- /dev/null +++ b/eval/src/tests/tensor/onnx_wrapper/int_types.py @@ -0,0 +1,33 @@ +# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +import onnx +from onnx import helper, TensorProto + +QUERY_TENSOR = helper.make_tensor_value_info('query_tensor', TensorProto.INT32, [1, 4]) +ATTRIBUTE_TENSOR = helper.make_tensor_value_info('attribute_tensor', TensorProto.INT32, [4, 1]) +BIAS_TENSOR = helper.make_tensor_value_info('bias_tensor', TensorProto.INT32, [1, 1]) +OUTPUT = helper.make_tensor_value_info('output', TensorProto.INT32, [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, + 'int_types_scoring', + [ + QUERY_TENSOR, + ATTRIBUTE_TENSOR, + BIAS_TENSOR, + ], + [OUTPUT], +) +model_def = helper.make_model(graph_def, producer_name='int_types.py') +onnx.save(model_def, 'int_types.onnx') |