Enable ONNX test in circle CI (#3144)

Summary: * Enable ONNX test in circle CI

Reviewed By: datumbox

Differential Revision: D25531035

fbshipit-source-id: 022b8613ac5bd8cb7165e9cbaa9daf99a9e1694e

.circleci/config.yml

@@ -145,6 +145,21 @@ jobs:
             pip install --user --progress-bar off --editable .
             python test/test_hub.py
 
+  torch_onnx_test:
+    docker:
+      - image: circleci/python:3.7
+    steps:
+      - checkout
+      - run:
+          command: |
+            pip install --user --progress-bar off numpy
+            pip install --user --progress-bar off --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+            # need to install torchvision dependencies due to transitive imports
+            pip install --user --progress-bar off --editable .
+            pip install --user onnx
+            pip install --user -i https://test.pypi.org/simple/ ort-nightly==1.5.2.dev202012031
+            python test/test_onnx.py
+
   binary_linux_wheel:
     <<: *binary_common
     docker:
@@ -1074,6 +1089,7 @@ workflows:
       - python_type_check
       - clang_format
       - torchhub_test
+      - torch_onnx_test
 
   unittest:
     jobs:
@@ -1189,6 +1205,7 @@ workflows:
       - python_type_check
       - clang_format
       - torchhub_test
+      - torch_onnx_test
       - binary_linux_wheel:
           cu_version: cpu
           filters:

.circleci/config.yml.in

@@ -145,6 +145,21 @@ jobs:
             pip install --user --progress-bar off --editable .
             python test/test_hub.py
 
+  torch_onnx_test:
+    docker:
+      - image: circleci/python:3.7
+    steps:
+      - checkout
+      - run:
+          command: |
+            pip install --user --progress-bar off numpy
+            pip install --user --progress-bar off --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+            # need to install torchvision dependencies due to transitive imports
+            pip install --user --progress-bar off --editable .
+            pip install --user onnx
+            pip install --user -i https://test.pypi.org/simple/ ort-nightly==1.5.2.dev202012031
+            python test/test_onnx.py
+
   binary_linux_wheel:
     <<: *binary_common
     docker:
@@ -700,6 +715,7 @@ workflows:
       - python_type_check
       - clang_format
       - torchhub_test
+      - torch_onnx_test
 
   unittest:
     jobs:
@@ -717,6 +733,7 @@ workflows:
       - python_type_check
       - clang_format
       - torchhub_test
+      - torch_onnx_test
       {{ build_workflows(prefix="nightly_", filter_branch="nightly", upload=True) }}
   docker_build:
     triggers:

test/test_onnx.py

@@ -170,7 +170,7 @@ def forward(self_module, images):
         input = torch.rand(3, 10, 20)
         input_test = torch.rand(3, 100, 150)
         self.run_model(TransformModule(), [(input,), (input_test,)],
-                       input_names=["input1"], dynamic_axes={"input1": [0, 1, 2, 3]})
+                       input_names=["input1"], dynamic_axes={"input1": [0, 1, 2]})
 
     def test_transform_images(self):
 
@@ -382,12 +382,12 @@ def test_faster_rcnn(self):
         # Test exported model on images of different size, or dummy input
         self.run_model(model, [(images,), (test_images,), (dummy_image,)], input_names=["images_tensors"],
                        output_names=["outputs"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "outputs": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2], "outputs": [0, 1, 2]},
                        tolerate_small_mismatch=True)
         # Test exported model for an image with no detections on other images
         self.run_model(model, [(dummy_image,), (images,)], input_names=["images_tensors"],
                        output_names=["outputs"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "outputs": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2], "outputs": [0, 1, 2]},
                        tolerate_small_mismatch=True)
 
     # Verify that paste_mask_in_image beahves the same in tracing.
@@ -434,16 +434,16 @@ def test_mask_rcnn(self):
         self.run_model(model, [(images,), (test_images,), (dummy_image,)],
                        input_names=["images_tensors"],
                        output_names=["boxes", "labels", "scores", "masks"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "boxes": [0, 1], "labels": [0],
-                                     "scores": [0], "masks": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2], "boxes": [0, 1], "labels": [0],
+                                     "scores": [0], "masks": [0, 1, 2]},
                        tolerate_small_mismatch=True)
         # TODO: enable this test once dynamic model export is fixed
         # Test exported model for an image with no detections on other images
         self.run_model(model, [(dummy_image,), (images,)],
                        input_names=["images_tensors"],
                        output_names=["boxes", "labels", "scores", "masks"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "boxes": [0, 1], "labels": [0],
-                                     "scores": [0], "masks": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2], "boxes": [0, 1], "labels": [0],
+                                     "scores": [0], "masks": [0, 1, 2]},
                        tolerate_small_mismatch=True)
 
     # Verify that heatmaps_to_keypoints behaves the same in tracing.
@@ -483,13 +483,13 @@ def test_keypoint_rcnn(self):
         self.run_model(model, [(images,), (test_images,), (dummy_images,)],
                        input_names=["images_tensors"],
                        output_names=["outputs1", "outputs2", "outputs3", "outputs4"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2]},
                        tolerate_small_mismatch=True)
 
         self.run_model(model, [(dummy_images,), (test_images,)],
                        input_names=["images_tensors"],
                        output_names=["outputs1", "outputs2", "outputs3", "outputs4"],
-                       dynamic_axes={"images_tensors": [0, 1, 2, 3]},
+                       dynamic_axes={"images_tensors": [0, 1, 2]},
                        tolerate_small_mismatch=True)