[ONNX] Fix ShuffleNetV2 model export issue. (#3158)

fatcat-z · datumbox · web-flow · commit 45d9a3040384 · 2020-12-12T11:30:07.000Z
* Fix an issue that ShuffleNetV2 model is exported to a wrong ONNX file if dynamic_axes field was provided.

* Add a ut for the bug fix.

* Fix flake8 issue.

* Don't access each element in x.shape, use x.size() instead.

Co-authored-by: Vasilis Vryniotis &lt;datumbox@users.noreply.github.com&gt;
diff --git a/test/test_onnx.py b/test/test_onnx.py
@@ -482,6 +482,17 @@ def test_keypoint_rcnn(self):
                        dynamic_axes={"images_tensors": [0, 1, 2]},
                        tolerate_small_mismatch=True)
 
+    def test_shufflenet_v2_dynamic_axes(self):
+        model = models.shufflenet_v2_x0_5(pretrained=True)
+        dummy_input = torch.randn(1, 3, 224, 224, requires_grad=True)
+        test_inputs = torch.cat([dummy_input, dummy_input, dummy_input], 0)
+
+        self.run_model(model, [(dummy_input,), (test_inputs,)],
+                       input_names=["input_images"],
+                       output_names=["output"],
+                       dynamic_axes={"input_images": {0: 'batch_size'}, "output": {0: 'batch_size'}},
+                       tolerate_small_mismatch=True)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py
@@ -19,7 +19,7 @@
 
 
 def channel_shuffle(x: Tensor, groups: int) -> Tensor:
-    batchsize, num_channels, height, width = x.data.size()
+    batchsize, num_channels, height, width = x.size()
     channels_per_group = num_channels // groups
 
     # reshape
