ONNX export for variable input sizes

test/test_onnx.py

@@ -28,14 +28,17 @@ class ONNXExporterTester(unittest.TestCase):
     def setUpClass(cls):
         torch.manual_seed(123)
 
-    def run_model(self, model, inputs_list, tolerate_small_mismatch=False, do_constant_folding=True):
+    def run_model(self, model, inputs_list, tolerate_small_mismatch=False, do_constant_folding=True, dynamic_axes=None,
+                  output_names=None, input_names=None):
         model.eval()
 
         onnx_io = io.BytesIO()
+        # onnx_io = '/home/neraoof/test/results/transform.onnx'
         # export to onnx with the first input
-        torch.onnx.export(model, inputs_list[0], onnx_io,
-                          do_constant_folding=do_constant_folding, opset_version=_onnx_opset_version)
 
+        torch.onnx.export(model, inputs_list[0], onnx_io,
+                          do_constant_folding=do_constant_folding, opset_version=_onnx_opset_version,
+                          dynamic_axes=dynamic_axes, input_names=input_names, output_names=output_names)
         # validate the exported model with onnx runtime
         for test_inputs in inputs_list:
             with torch.no_grad():
@@ -65,6 +68,7 @@ def to_numpy(tensor):
         # compute onnxruntime output prediction
         ort_inputs = dict((ort_session.get_inputs()[i].name, inpt) for i, inpt in enumerate(inputs))
         ort_outs = ort_session.run(None, ort_inputs)
+
         for i in range(0, len(outputs)):
             try:
                 torch.testing.assert_allclose(outputs[i], ort_outs[i], rtol=1e-03, atol=1e-05)
@@ -123,9 +127,9 @@ def __init__(self_module):
             def forward(self_module, images):
                 return self_module.transform(images)[0].tensors
 
-        input = [torch.rand(3, 100, 200), torch.rand(3, 200, 200)]
-        input_test = [torch.rand(3, 100, 200), torch.rand(3, 200, 200)]
-        self.run_model(TransformModule(), [input, input_test])
+        input = torch.rand(3, 100, 200), torch.rand(3, 200, 200)
+        input_test = torch.rand(3, 100, 200), torch.rand(3, 200, 200)
+        self.run_model(TransformModule(), [(input,), (input_test,)])
 
     def _init_test_generalized_rcnn_transform(self):
         min_size = 100
@@ -207,22 +211,28 @@ def get_features(self, images):
 
     def test_rpn(self):
         class RPNModule(torch.nn.Module):
-            def __init__(self_module, images):
+            def __init__(self_module):
                 super(RPNModule, self_module).__init__()
                 self_module.rpn = self._init_test_rpn()
-                self_module.images = ImageList(images, [i.shape[-2:] for i in images])
 
-            def forward(self_module, features):
-                return self_module.rpn(self_module.images, features)
+            def forward(self_module, images, features):
+                images = ImageList(images, [i.shape[-2:] for i in images])
+                return self_module.rpn(images, features)
 
         images = torch.rand(2, 3, 600, 600)
         features = self.get_features(images)
-        test_features = self.get_features(images)
+        images2 = torch.rand(2, 3, 1000, 1000)
+        test_features = self.get_features(images2)
 
-        model = RPNModule(images)
+        model = RPNModule()
         model.eval()
-        model(features)
-        self.run_model(model, [(features,), (test_features,)], tolerate_small_mismatch=True)
+        model(images, features)
+
+        self.run_model(model, [(images, features), (images2, test_features)], tolerate_small_mismatch=True,
+                       input_names=["input1", "input2", "input3", "input4", "input5", "input6"],
+                       dynamic_axes={"input1": [0, 1, 2, 3], "input2": [0, 1, 2, 3],
+                                     "input3": [0, 1, 2, 3], "input4": [0, 1, 2, 3],
+                                     "input5": [0, 1, 2, 3], "input6": [0, 1, 2, 3]})
 
     def test_multi_scale_roi_align(self):
 
@@ -251,32 +261,37 @@ def forward(self, input, boxes):
 
     def test_roi_heads(self):
         class RoiHeadsModule(torch.nn.Module):
-            def __init__(self_module, images):
+            def __init__(self_module):
                 super(RoiHeadsModule, self_module).__init__()
                 self_module.transform = self._init_test_generalized_rcnn_transform()
                 self_module.rpn = self._init_test_rpn()
                 self_module.roi_heads = self._init_test_roi_heads_faster_rcnn()
-                self_module.original_image_sizes = [img.shape[-2:] for img in images]
-                self_module.images = ImageList(images, [i.shape[-2:] for i in images])
 
-            def forward(self_module, features):
-                proposals, _ = self_module.rpn(self_module.images, features)
-                detections, _ = self_module.roi_heads(features, proposals, self_module.images.image_sizes)
+            def forward(self_module, images, features):
+                original_image_sizes = [img.shape[-2:] for img in images]
+                images = ImageList(images, [i.shape[-2:] for i in images])
+                proposals, _ = self_module.rpn(images, features)
+                detections, _ = self_module.roi_heads(features, proposals, images.image_sizes)
                 detections = self_module.transform.postprocess(detections,
-                                                               self_module.images.image_sizes,
-                                                               self_module.original_image_sizes)
+                                                               images.image_sizes,
+                                                               original_image_sizes)
                 return detections
 
         images = torch.rand(2, 3, 600, 600)
         features = self.get_features(images)
-        test_features = self.get_features(images)
+        images2 = torch.rand(2, 3, 1000, 1000)
+        test_features = self.get_features(images2)
 
-        model = RoiHeadsModule(images)
+        model = RoiHeadsModule()
         model.eval()
-        model(features)
-        self.run_model(model, [(features,), (test_features,)])
+        model(images, features)
+
+        self.run_model(model, [(images, features), (images2, test_features)], tolerate_small_mismatch=True,
+                       input_names=["input1", "input2", "input3", "input4", "input5", "input6"],
+                       dynamic_axes={"input1": [0, 1, 2, 3], "input2": [0, 1, 2, 3], "input3": [0, 1, 2, 3],
+                                     "input4": [0, 1, 2, 3], "input5": [0, 1, 2, 3], "input6": [0, 1, 2, 3]})
 
-    def get_image_from_url(self, url):
+    def get_image_from_url(self, url, size=None):
         import requests
         import numpy
         from PIL import Image
@@ -285,29 +300,35 @@ def get_image_from_url(self, url):
 
         data = requests.get(url)
         image = Image.open(BytesIO(data.content)).convert("RGB")
-        image = image.resize((300, 200), Image.BILINEAR)
+
+        if size is None:
+            size = (300, 200)
+        image = image.resize(size, Image.BILINEAR)
 
         to_tensor = transforms.ToTensor()
         return to_tensor(image)
 
     def get_test_images(self):
         image_url = "http://farm3.staticflickr.com/2469/3915380994_2e611b1779_z.jpg"
-        image = self.get_image_from_url(url=image_url)
+        image = self.get_image_from_url(url=image_url, size=(800, 1201))
+
         image_url2 = "https://pytorch.org/tutorials/_static/img/tv_tutorial/tv_image05.png"
-        image2 = self.get_image_from_url(url=image_url2)
+        image2 = self.get_image_from_url(url=image_url2, size=(873, 800))
+
         images = [image]
         test_images = [image2]
         return images, test_images
 
     def test_faster_rcnn(self):
         images, test_images = self.get_test_images()
 
-        model = models.detection.faster_rcnn.fasterrcnn_resnet50_fpn(pretrained=True,
-                                                                     min_size=200,
-                                                                     max_size=300)
+        model = models.detection.faster_rcnn.fasterrcnn_resnet50_fpn(pretrained=True)
         model.eval()
         model(images)
-        self.run_model(model, [(images,), (test_images,)])
+        self.run_model(model, [(images,), (test_images,)], input_names=["images_tensors"],
+                       output_names=["outputs"],
+                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "outputs": [0, 1, 2, 3]},
+                       tolerate_small_mismatch=True)
 
     # Verify that paste_mask_in_image beahves the same in tracing.
     # This test also compares both paste_masks_in_image and _onnx_paste_masks_in_image
@@ -346,10 +367,14 @@ def test_paste_mask_in_image(self):
     def test_mask_rcnn(self):
         images, test_images = self.get_test_images()
 
-        model = models.detection.mask_rcnn.maskrcnn_resnet50_fpn(pretrained=True, min_size=200, max_size=300)
+        model = models.detection.mask_rcnn.maskrcnn_resnet50_fpn(pretrained=True)
         model.eval()
         model(images)
-        self.run_model(model, [(images,), (test_images,)])
+        self.run_model(model, [(images,), (test_images,)],
+                       input_names=["images_tensors"],
+                       output_names=["outputs"],
+                       dynamic_axes={"images_tensors": [0, 1, 2, 3], "outputs": [0, 1, 2, 3]},
+                       tolerate_small_mismatch=True)
 
     # Verify that heatmaps_to_keypoints behaves the same in tracing.
     # This test also compares both heatmaps_to_keypoints and _onnx_heatmaps_to_keypoints
@@ -383,9 +408,7 @@ def test_keypoint_rcnn(self):
         class KeyPointRCNN(torch.nn.Module):
             def __init__(self):
                 super(KeyPointRCNN, self).__init__()
-                self.model = models.detection.keypoint_rcnn.keypointrcnn_resnet50_fpn(pretrained=True,
-                                                                                      min_size=200,
-                                                                                      max_size=300)
+                self.model = models.detection.keypoint_rcnn.keypointrcnn_resnet50_fpn(pretrained=True)
 
             def forward(self, images):
                 output = self.model(images)
@@ -398,7 +421,11 @@ def forward(self, images):
         model = KeyPointRCNN()
         model.eval()
         model(test_images)
-        self.run_model(model, [(images,), (test_images,)])
+        self.run_model(model, [(images,), (test_images,)],
+                       input_names=["images_tensors"],
+                       output_names=["outputs1", "outputs2", "outputs3", "outputs4"],
+                       dynamic_axes={"images_tensors": [0, 1, 2, 3]},
+                       tolerate_small_mismatch=True)
 
 
 if __name__ == '__main__':

torchvision/models/detection/rpn.py

@@ -17,8 +17,7 @@
 @torch.jit.unused
 def _onnx_get_num_anchors_and_pre_nms_top_n(ob, orig_pre_nms_top_n):
     # type: (Tensor, int) -> Tuple[int, int]
-    from torch.onnx import operators
-    num_anchors = operators.shape_as_tensor(ob)[1].unsqueeze(0)
+    num_anchors = ob.shape[1].unsqueeze(0)
     # TODO : remove cast to IntTensor/num_anchors.dtype when
     #        ONNX Runtime version is updated with ReduceMin int64 support
     pre_nms_top_n = torch.min(torch.cat(
@@ -157,7 +156,8 @@ def forward(self, image_list, feature_maps):
         # type: (ImageList, List[Tensor])
         grid_sizes = list([feature_map.shape[-2:] for feature_map in feature_maps])
         image_size = image_list.tensors.shape[-2:]
-        strides = [[int(image_size[0] / g[0]), int(image_size[1] / g[1])] for g in grid_sizes]
+        strides = [[torch.tensor(image_size[0] / g[0], dtype=torch.int64),
+                    torch.tensor(image_size[1] / g[1], dtype=torch.int64)] for g in grid_sizes]
         dtype, device = feature_maps[0].dtype, feature_maps[0].device
         self.set_cell_anchors(dtype, device)
         anchors_over_all_feature_maps = self.cached_grid_anchors(grid_sizes, strides)
@@ -354,7 +354,17 @@ def _get_top_n_idx(self, objectness, num_anchors_per_level):
         # type: (Tensor, List[int])
         r = []
         offset = 0
-        for ob in objectness.split(num_anchors_per_level, 1):
+        if torchvision._is_tracing():
+            # Split's split_size is traced as constant in onnx exporting, use Gather
+            start_list = [torch.tensor(0)]
+            end_list = [num_anchors_per_level[0].clone()]
+            for cnt in num_anchors_per_level[1:]:
+                start_list.append(end_list[-1].clone())
+                end_list.append(end_list[-1] + cnt)
+            objectness_per_level = [objectness[:, s:e] for s, e in zip(start_list, end_list)]
+        else:
+            objectness_per_level = objectness.split(num_anchors_per_level, 1)
+        for ob in objectness_per_level:
             if torchvision._is_tracing():
                 num_anchors, pre_nms_top_n = _onnx_get_num_anchors_and_pre_nms_top_n(ob, self.pre_nms_top_n())
             else:
@@ -372,14 +382,12 @@ def filter_proposals(self, proposals, objectness, image_shapes, num_anchors_per_
         # do not backprop throught objectness
         objectness = objectness.detach()
         objectness = objectness.reshape(num_images, -1)
-
         levels = [
             torch.full((n,), idx, dtype=torch.int64, device=device)
             for idx, n in enumerate(num_anchors_per_level)
         ]
         levels = torch.cat(levels, 0)
         levels = levels.reshape(1, -1).expand_as(objectness)
-
         # select top_n boxes independently per level before applying nms
         top_n_idx = self._get_top_n_idx(objectness, num_anchors_per_level)
 
@@ -466,7 +474,9 @@ def forward(self, images, features, targets=None):
         anchors = self.anchor_generator(images, features)
 
         num_images = len(anchors)
-        num_anchors_per_level = [o[0].numel() for o in objectness]
+        num_anchors_per_level_shape_tensors = [o[0].shape for o in objectness]
+        num_anchors_per_level = [s[0] * s[1] * s[2] for s in num_anchors_per_level_shape_tensors]
+
         objectness, pred_bbox_deltas = \
             concat_box_prediction_layers(objectness, pred_bbox_deltas)
         # apply pred_bbox_deltas to anchors to obtain the decoded proposals

torchvision/models/detection/transform.py

@@ -90,7 +90,8 @@ def resize(self, image, target):
         if max_size * scale_factor > self.max_size:
             scale_factor = self.max_size / max_size
         image = torch.nn.functional.interpolate(
-            image[None], scale_factor=scale_factor, mode='bilinear', align_corners=False)[0]
+            image[None], scale_factor=scale_factor, mode='bilinear',
+            align_corners=False, recompute_scale_factor=False)[0]
 
         if target is None:
             return image, target

torchvision/ops/boxes.py

@@ -3,6 +3,7 @@
 import torch
 from torch.jit.annotations import Tuple
 from torch import Tensor
+import torchvision
 
 
 def nms(boxes, scores, iou_threshold):
@@ -112,8 +113,19 @@ def clip_boxes_to_image(boxes, size):
     boxes_x = boxes[..., 0::2]
     boxes_y = boxes[..., 1::2]
     height, width = size
-    boxes_x = boxes_x.clamp(min=0, max=width)
-    boxes_y = boxes_y.clamp(min=0, max=height)
+
+    if torchvision._is_tracing():
+        height = height.to(torch.float32)
+        width = width.to(torch.float32)
+
+        boxes_x = torch.max(boxes_x, torch.tensor(0.))
+        boxes_x = torch.min(boxes_x, width)
+        boxes_y = torch.max(boxes_y, torch.tensor(0.))
+        boxes_y = torch.min(boxes_y, height)
+    else:
+        boxes_x = boxes_x.clamp(min=0, max=width)
+        boxes_y = boxes_y.clamp(min=0, max=height)
+
     clipped_boxes = torch.stack((boxes_x, boxes_y), dim=dim)
     return clipped_boxes.reshape(boxes.shape)