Add unit test for PointRend ONNX export

Refer to projects/PointRend/configs/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco.yaml
facebookresearch · Jan 19, 2023 · 0dfcbb5 · 0dfcbb5
1 parent 650c60e
commit 0dfcbb5
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/test_export_onnx.py b/tests/test_export_onnx.py
@@ -88,6 +88,19 @@ def inference_func(model, image):
             opset_version=STABLE_ONNX_OPSET_VERSION,
         )
 
+    @skipIfUnsupportedMinOpsetVersion(16, STABLE_ONNX_OPSET_VERSION)
+    @skipIfUnsupportedMinTorchVersion("1.11.1")
+    def test_pointrend_rcnn_R_50_FPN_3x_coco(self):
+        def inference_func(model, image):
+            inputs = [{"image": image, "height": image.shape[1], "width": image.shape[2]}]
+            return model.inference(inputs, do_postprocess=True)[0]["instances"]
+
+        self._test_model_from_config_path(
+            "projects/PointRend/configs/InstanceSegmentation/pointrend_rcnn_R_50_FPN_3x_coco.yaml",
+            inference_func,
+            opset_version=STABLE_ONNX_OPSET_VERSION,
+        )
+
     ################################################################################
     # Testcase internals - DO NOT add tests below this point
     ################################################################################