Fix use of ONNX optimizer by Caffe2 backend

Thiago Crepaldi · Thiago Crepaldi · commit 83f6441be700 · 2022-04-13T06:10:28.000-07:00
diff --git a/caffe2/python/onnx/backend.py b/caffe2/python/onnx/backend.py
@@ -651,7 +651,13 @@ def optimize_onnx(input, init=False, predict=False):
             passes.append('split_init')
         if predict:
             passes.append('split_predict')
-        out = onnx.optimizer.optimize(input, passes)
+        try:
+            out = onnx.optimizer.optimize(input, passes)
+        except AttributeError:
+            warnings.warn("OptimizerWarning: optimizer module not found in ONNX version {}".format(onnx.__version__))
+            # ONNX does no ship onnx.optimizer since version 1.9+
+            import onnxoptimizer
+            out = onnxoptimizer.optimize(input, passes)
         return out
 
     @classmethod
@@ -881,8 +887,9 @@ def _onnx_model_to_caffe2_net(cls, onnx_model, device, opset_version, include_in
         try:
             init_model = cls.optimize_onnx(onnx_model, init=True)
             pred_model = cls.optimize_onnx(onnx_model, predict=True)
-        except AttributeError:
-            warnings.warn("OptimizerWarning: optimizer module not found in ONNX version {}".format(onnx.__version__))
+        except ModuleNotFoundError:
+            warnings.warn("OptimizerWarning: onnxoptimizer module not installed. "
+                          "init_model and pred_model models will not be splitted, which can cause a runtime error")
             init_model = onnx_model
             pred_model = onnx_model
 
