fix to emit appropriate import error

niboshi · Oct 14, 2019 · bd4613f · bd4613f
1 parent 5939fb3
commit bd4613f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/onnx_chainer/export.py b/onnx_chainer/export.py
@@ -4,19 +4,19 @@
 import warnings
 
 import chainer
-import onnx
-from onnx.mapping import NP_TYPE_TO_TENSOR_TYPE
-from onnx import shape_inference
-
-from onnx_chainer.context import Context
-from onnx_chainer.graph import Graph
-from onnx_chainer import mapping
-from onnx_chainer.onnx_helper import is_support_non_standard_domain
 
 try:
+    import onnx
     from onnx import checker
     from onnx import helper
+    from onnx.mapping import NP_TYPE_TO_TENSOR_TYPE
     from onnx import numpy_helper
+    from onnx import shape_inference
+
+    from onnx_chainer.context import Context
+    from onnx_chainer.graph import Graph
+    from onnx_chainer import mapping
+    from onnx_chainer.onnx_helper import is_support_non_standard_domain
 
     _available = True
 except ImportError:
@@ -31,7 +31,7 @@ def _check_available():
         raise ImportError(
             'ONNX is not installed on your environment. Exporting your model '
             'in ONNX format needs the onnx package.\n\n'
-            '\t$ pip install onnx\n\n')
+            '\t$ pip install \'onnx<1.6.0\'\n\n')
 
 
 def convert_parameter(parameter, context):

diff --git a/onnx_chainer/export_testcase.py b/onnx_chainer/export_testcase.py
@@ -3,9 +3,13 @@
 
 import chainer
 
-from onnx_chainer.export import export
-from onnx_chainer.onnx_helper import cleanse_param_name
-from onnx_chainer.onnx_helper import write_tensor_pb
+from onnx_chainer.export import _available
+
+if _available:
+    from onnx_chainer.export import export
+
+    from onnx_chainer.onnx_helper import cleanse_param_name
+    from onnx_chainer.onnx_helper import write_tensor_pb
 
 
 def export_testcase(model, args, out_dir, output_grad=False, **kwargs):