From 6d1351f298c1d44ee012aac7be2ad89306d9ef80 Mon Sep 17 00:00:00 2001 From: bddppq Date: Fri, 28 Jun 2019 11:04:49 -0400 Subject: [PATCH] Expose ONNX_ML build option to python (#2138) * Expose ONNX_ML build option to python * Add type --- onnx/__init__.py | 1 + onnx/cpp2py_export.cc | 8 ++++++++ onnx/onnx_cpp2py_export/__init__.pyi | 2 ++ onnx/test/shape_inference_test.py | 10 +++------- 4 files changed, 14 insertions(+), 7 deletions(-) diff --git a/onnx/__init__.py b/onnx/__init__.py index ab69280dde4..b9ca9620a94 100644 --- a/onnx/__init__.py +++ b/onnx/__init__.py @@ -5,6 +5,7 @@ import os +from .onnx_cpp2py_export import ONNX_ML from onnx.external_data_helper import load_external_data_for_model, write_external_data_tensors from .onnx_pb import * # noqa from .onnx_operators_pb import * # noqa diff --git a/onnx/cpp2py_export.cc b/onnx/cpp2py_export.cc index c4021f9c48b..c04267c7f8a 100644 --- a/onnx/cpp2py_export.cc +++ b/onnx/cpp2py_export.cc @@ -19,6 +19,14 @@ using namespace pybind11::literals; PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) { onnx_cpp2py_export.doc() = "Python interface to onnx"; + onnx_cpp2py_export.attr("ONNX_ML") = py::bool_( +#ifdef ONNX_ML + true +#else // ONNX_ML + false +#endif // ONNX_ML + ); + // Submodule `schema` auto defs = onnx_cpp2py_export.def_submodule("defs"); defs.doc() = "Schema submodule"; diff --git a/onnx/onnx_cpp2py_export/__init__.pyi b/onnx/onnx_cpp2py_export/__init__.pyi index 423ac367ffc..232bc303872 100644 --- a/onnx/onnx_cpp2py_export/__init__.pyi +++ b/onnx/onnx_cpp2py_export/__init__.pyi @@ -1,3 +1,5 @@ # This is __init__.pyi, not __init__.py # This directory is only considered for typing information, not for actual # module implementations. + +ONNX_ML : bool = ... diff --git a/onnx/test/shape_inference_test.py b/onnx/test/shape_inference_test.py index 58fe69dc696..5d6b84fa6df 100644 --- a/onnx/test/shape_inference_test.py +++ b/onnx/test/shape_inference_test.py @@ -3,7 +3,7 @@ from __future__ import print_function from __future__ import unicode_literals -from onnx import checker, helper, TensorProto, NodeProto, GraphProto, ValueInfoProto, ModelProto +from onnx import checker, helper, TensorProto, NodeProto, GraphProto, ValueInfoProto, ModelProto, ONNX_ML from onnx.helper import make_node, make_tensor, make_tensor_value_info, make_empty_tensor_value_info, make_opsetid from typing import Sequence, Union, Text, Tuple, List, Any, Optional import onnx.shape_inference @@ -1943,9 +1943,7 @@ def test_tile_rank_inference(self): # type: () -> None self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.FLOAT, (None, None, None))]) # type: ignore def test_linearclassifier_1D_input(self): # type: () -> None - onnx_ml = os.environ.get('ONNX_ML') # type: ignore - # No environment variable set (None) indicates ONNX_ML=1 - if(onnx_ml is None or int(onnx_ml) != 0): + if ONNX_ML: graph = self._make_graph( [('x', TensorProto.FLOAT, (5,))], [make_node('LinearClassifier', ['x'], ['y', 'z'], domain='ai.onnx.ml', coefficients=[0.0008, -0.0008], intercepts=[2.0, 2.0], classlabels_ints=[1, 2])], @@ -1955,9 +1953,7 @@ def test_linearclassifier_1D_input(self): # type: () -> None opset_imports=[make_opsetid('ai.onnx.ml', 1), make_opsetid('', 11)]) def test_linearclassifier_2D_input(self): # type: () -> None - onnx_ml = os.environ.get('ONNX_ML') # type: ignore - # No environment variable set (None) indicates ONNX_ML=1 - if(onnx_ml is None or int(onnx_ml) != 0): + if ONNX_ML: graph = self._make_graph( [('x', TensorProto.FLOAT, (4, 5))], [make_node('LinearClassifier', ['x'], ['y', 'z'], domain='ai.onnx.ml', coefficients=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6], intercepts=[2.0, 2.0, 3.0], classlabels_ints=[1, 2, 3])],