Fix first reshape in pixel_shuffle conversion (#21486)

Kevin Chen · facebook-github-bot · commit f87d5cc191b5 · 2019-06-13T11:44:54.000-07:00
Summary: When converting pixel_shuffle to reshape + transpose + reshape, the first reshape should be: [N, C * r^2, H, W] => [N, C, r, r, H, W] in order to match pytorch's implementation (see ATen PixelShuffle.cpp). This previously wasn't caught by the test case, since it uses C = r = 4. Updated test case to have C = 2, r = 4. Pull Request resolved: #21486 Reviewed By: houseroad Differential Revision: D15700945 Pulled By: houseroad fbshipit-source-id: 47019691fdc20e152e867c7f6fd57da104a12948
diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py
@@ -1112,7 +1112,7 @@ def test_instance_norm(self):
 
     def test_pixel_shuffle(self):
         underlying = nn.PixelShuffle(4)
-        shape = (1, 64, 5, 5)
+        shape = (1, 32, 5, 5)
         input = Variable(torch.randn(*shape),
                          requires_grad=True)
         self.run_model_test(underlying, train=False, input=(input),
diff --git a/torch/onnx/symbolic_opset9.py b/torch/onnx/symbolic_opset9.py
@@ -1248,8 +1248,8 @@ def pixel_shuffle(g, self, upscale_factor):
     if len(dims) != 4:
         return _unimplemented("pixel_shuffle", "only support 4d input")
     output_channel = dims[1] // upscale_factor // upscale_factor
-    after_view = view(g, self, [-1, upscale_factor, upscale_factor,
-                                output_channel, dims[2], dims[3]])
+    after_view = view(g, self, [-1, output_channel, upscale_factor, upscale_factor,
+                                dims[2], dims[3]])
     after_transpose = g.op("Transpose", after_view, perm_i=[0, 1, 4, 2, 5, 3])
     return view(g, after_transpose,
                 [-1, output_channel, dims[2] * upscale_factor, dims[3] *
