From b5a6303c5182d4bb503790eb2a76a8e0f0dbd8bf Mon Sep 17 00:00:00 2001 From: Jay Zhang Date: Fri, 11 Dec 2020 14:38:34 +0800 Subject: [PATCH 1/4] Fix an issue that ShuffleNetV2 model is exported to a wrong ONNX file if dynamic_axes field was provided. --- torchvision/models/shufflenetv2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index 9ba090ad09b..54f54e01ef2 100644 --- 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.shape[0], x.shape[1], x.shape[2], x.shape[3] channels_per_group = num_channels // groups # reshape From f1124ee852cc955777988c4dd40eba67a0b48aaa Mon Sep 17 00:00:00 2001 From: Jay Zhang Date: Fri, 11 Dec 2020 19:18:25 +0800 Subject: [PATCH 2/4] Add a ut for the bug fix. --- test/test_onnx.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/test_onnx.py b/test/test_onnx.py index 54a4e385a8d..0da3bc620b5 100644 --- a/test/test_onnx.py +++ b/test/test_onnx.py @@ -1,4 +1,5 @@ import io +import numpy as np import torch from torchvision import ops from torchvision import models @@ -478,6 +479,16 @@ 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() From d8dd6c2968687ccbcbabb2845b1d4eaec031bdaf Mon Sep 17 00:00:00 2001 From: Jay Zhang Date: Fri, 11 Dec 2020 19:25:05 +0800 Subject: [PATCH 3/4] Fix flake8 issue. --- test/test_onnx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_onnx.py b/test/test_onnx.py index 0da3bc620b5..866e28f59a5 100644 --- a/test/test_onnx.py +++ b/test/test_onnx.py @@ -1,5 +1,4 @@ import io -import numpy as np import torch from torchvision import ops from torchvision import models @@ -490,5 +489,6 @@ def test_shufflenet_v2_dynamic_axes(self): dynamic_axes={"input_images": {0: 'batch_size'}, "output": {0: 'batch_size'}}, tolerate_small_mismatch=True) + if __name__ == '__main__': unittest.main() From 5c187df2a2441867bc0b00a2c31a22049a9e6f8b Mon Sep 17 00:00:00 2001 From: Jay Zhang Date: Sat, 12 Dec 2020 13:08:35 +0800 Subject: [PATCH 4/4] Don't access each element in x.shape, use x.size() instead. --- torchvision/models/shufflenetv2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index 54f54e01ef2..9a4333eb10b 100644 --- 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.shape[0], x.shape[1], x.shape[2], x.shape[3] + batchsize, num_channels, height, width = x.size() channels_per_group = num_channels // groups # reshape