add NCHW2NHWC and NHWC2NCHW in utils.py (pytorch#15588)

Summary:
Pull Request resolved: pytorch#15588

Use NHWC2NCHW or NCHW2NHWC functions which is easier to understand compared to code using transpose and generalizable to non-2D convolutions.

Reviewed By: csummersea

Differential Revision: D13557674

fbshipit-source-id: c4fdb8850503ea58f6b17b188513ae2b29691ec0

caffe2/python/operator_test/channel_shuffle_test.py

@@ -3,7 +3,7 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from caffe2.python import core
+from caffe2.python import core, utils
 import caffe2.python.hypothesis_test_util as hu
 import caffe2.python.serialized_test.serialized_test_util as serial
 from hypothesis import given

caffe2/python/operator_test/conv_test.py

@@ -10,7 +10,7 @@
 import hypothesis.strategies as st
 
 from caffe2.proto import caffe2_pb2
-from caffe2.python import brew, core, workspace
+from caffe2.python import brew, core, utils, workspace
 import caffe2.python.hip_test_util as hiputl
 import caffe2.python.hypothesis_test_util as hu
 from caffe2.python.model_helper import ModelHelper
@@ -56,14 +56,6 @@ def _cudnn_convolution_algo_count(direction):
         return st.sampled_from([-1])
 
 
-def nhwc2nchw(tensor):
-    return tensor.transpose((0, tensor.ndim - 1) + tuple(range(1, tensor.ndim - 1)))
-
-
-def nchw2nhwc(tensor):
-    return tensor.transpose((0,) + tuple(range(2, tensor.ndim)) + (1,))
-
-
 class TestConvolution(serial.SerializedTestCase):
     # CUDNN does NOT support different padding values and we skip it
     @given(op_type=st.sampled_from(["Conv", "Conv2D"]),
@@ -119,8 +111,8 @@ def test_convolution_separate_stride_pad_gradients(
         ).astype(np.float32) - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, w, b] if use_bias else [X, w]
 
@@ -178,8 +170,8 @@ def test_convolution_separate_stride_pad_layout(
                 device_option=gc,
             )
             if order == "NCHW":
-                X_f = X.transpose((0, 3, 1, 2))
-                w_f = w.transpose((0, 3, 1, 2))
+                X_f = utils.NHWC2NCHW(X)
+                w_f = utils.NHWC2NCHW(w)
             else:
                 X_f = X
                 w_f = w
@@ -190,7 +182,7 @@ def test_convolution_separate_stride_pad_layout(
             outputs[order] = self.ws.blobs["Y"].fetch()
         np.testing.assert_allclose(
             outputs["NCHW"],
-            outputs["NHWC"].transpose((0, 3, 1, 2)),
+            utils.NHWC2NCHW(outputs["NHWC"]),
             atol=1e-4,
             rtol=1e-4)
 
@@ -260,8 +252,8 @@ def test_convolution_gradients(
         ).astype(np.float32) - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, w, b] if use_bias else [X, w]
         # Error handling path.
@@ -320,8 +312,8 @@ def _nd_convolution(self, n, input_channels, output_channels,
             w = np.random.rand(*filter_dims).astype(np.float32) - 0.5
             b = np.random.rand(output_channels).astype(np.float32) - 0.5
             if order == "NHWC":
-                X = nchw2nhwc(X)
-                w = nchw2nhwc(w)
+                X = utils.NCHW2NHWC(X)
+                w = utils.NCHW2NHWC(w)
 
             inputs = [X, w, b] if use_bias else [X, w]
 
@@ -523,8 +515,8 @@ def test_convolution_layout(self, op_type, stride, pad, kernel, dilation,
                     exhaustive_search=True,
                 )
                 if order == "NCHW":
-                    X_f = X.transpose((0, 3, 1, 2))
-                    w_f = w.transpose((0, 3, 1, 2))
+                    X_f = utils.NHWC2NCHW(X)
+                    w_f = utils.NHWC2NCHW(w)
                 else:
                     X_f = X
                     w_f = w
@@ -542,7 +534,7 @@ def test_convolution_layout(self, op_type, stride, pad, kernel, dilation,
 
         def canonical(o):
             if o.order == "NHWC":
-                return o.Y.transpose((0, 3, 1, 2))
+                return utils.NHWC2NCHW(o.Y)
             else:
                 return o.Y
 

caffe2/python/operator_test/conv_transpose_test.py

@@ -6,7 +6,7 @@
 from hypothesis import assume, given, settings
 import hypothesis.strategies as st
 
-from caffe2.python import core
+from caffe2.python import core, utils
 import caffe2.python.hypothesis_test_util as hu
 import caffe2.python.hip_test_util as hiputl
 
@@ -57,8 +57,8 @@ def test_convolution_transpose_layout_legacy_args(
                 device_option=gc,
             )
             if order == "NCHW":
-                X_f = X.transpose((0, 3, 1, 2))
-                w_f = w.transpose((0, 3, 1, 2))
+                X_f = utils.NHWC2NCHW(X)
+                w_f = utils.NHWC2NCHW(w)
             else:
                 X_f = X
                 w_f = w
@@ -78,7 +78,7 @@ def test_convolution_transpose_layout_legacy_args(
             (batch_size, output_channels, output_size, output_size))
         np.testing.assert_allclose(
             outputs["NCHW"],
-            outputs["NHWC"].transpose((0, 3, 1, 2)),
+            utils.NHWC2NCHW(outputs["NHWC"]),
             atol=1e-4,
             rtol=1e-4)
 
@@ -127,8 +127,8 @@ def test_convolution_transpose_layout(
                 device_option=gc,
             )
             if order == "NCHW":
-                X_f = X.transpose((0, 3, 1, 2))
-                w_f = w.transpose((0, 3, 1, 2))
+                X_f = utils.NHWC2NCHW(X)
+                w_f = utils.NHWC2NCHW(w)
             else:
                 X_f = X
                 w_f = w
@@ -148,7 +148,7 @@ def test_convolution_transpose_layout(
             (batch_size, output_channels, output_size, output_size))
         np.testing.assert_allclose(
             outputs["NCHW"],
-            outputs["NHWC"].transpose((0, 3, 1, 2)),
+            utils.NHWC2NCHW(outputs["NHWC"]),
             atol=1e-4,
             rtol=1e-4)
 
@@ -201,8 +201,8 @@ def test_convolution_transpose_separate_stride_pad_adj_layout(
                 device_option=gc,
             )
             if order == "NCHW":
-                X_f = X.transpose((0, 3, 1, 2))
-                w_f = w.transpose((0, 3, 1, 2))
+                X_f = utils.NHWC2NCHW(X)
+                w_f = utils.NHWC2NCHW(w)
             else:
                 X_f = X
                 w_f = w
@@ -223,7 +223,7 @@ def test_convolution_transpose_separate_stride_pad_adj_layout(
             (batch_size, output_channels, output_h, output_w))
         np.testing.assert_allclose(
             outputs["NCHW"],
-            outputs["NHWC"].transpose((0, 3, 1, 2)),
+            utils.NHWC2NCHW(outputs["NHWC"]),
             atol=1e-4,
             rtol=1e-4)
 
@@ -268,8 +268,8 @@ def test_convolution_transpose_gradients(self, stride, pad, kernel, adj,
             no_gradient_to_input=not compute_dX,
         )
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, w, b] if use_bias else [X, w]
         self.assertDeviceChecks(dc, op, inputs, [0])
@@ -339,8 +339,8 @@ def test_convolution_transpose_separate_stride_pad_adj_gradient(
             no_gradient_to_input=not compute_dX,
         )
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, w, b] if use_bias else [X, w]
         self.assertDeviceChecks(dc, op, inputs, [0])

caffe2/python/operator_test/deform_conv_test.py

@@ -7,13 +7,11 @@
 import hypothesis.strategies as st
 
 from caffe2.proto import caffe2_pb2
-from caffe2.python import core, workspace
+from caffe2.python import core, utils, workspace
 import caffe2.python.hypothesis_test_util as hu
 import unittest
 import os
 
-import unittest
-
 
 def _cudnn_supports(
         dilation=False,
@@ -108,7 +106,7 @@ def _conv_2d_shuffle_offsets(
     w0 = [[w0] * input_channels] * output_channels
     return (
         np.array([e] * batch_size).astype(np.float32),
-        np.array(w0).astype(np.float32).transpose((0, 2, 3, 1))
+        utils.NCHW2NHWC(np.array(w0).astype(np.float32))
     )
 
 
@@ -163,8 +161,8 @@ def test_null_offset_convolution(self, stride, pad, kernel, dilation, size,
             - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, o, w, b] if use_bias else [X, o, w]
 
@@ -247,8 +245,8 @@ def test_flat_input_convolution(self, stride, pad, kernel, dilation, size,
         w = np.ones((output_channels, kernel, kernel, input_channels), np.float32) - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, o, w, b] if use_bias else [X, o, w]
 
@@ -334,9 +332,9 @@ def test_shuffle_input_convolution(self, stride, pad, kernel, dilation, size,
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
 
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
-            w0 = w0.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
+            w0 = utils.NHWC2NCHW(w0)
 
         inputs = [X, o, w, b] if use_bias else [X, o, w]
 
@@ -425,8 +423,8 @@ def test_conv_separate_stride_pad_gradients(self, stride_h, stride_w,
             - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, o, w, b] if use_bias else [X, o, w]
 
@@ -495,8 +493,8 @@ def test_conv_gradients(self, stride, pad, kernel, dilation, size,
             - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, o, w, b] if use_bias else [X, o, w]
         # Error handling path.

caffe2/python/operator_test/depthwise_3x3_conv_test.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 import caffe2.python.hypothesis_test_util as hu
-from caffe2.python import core, dyndep, workspace
+from caffe2.python import core, dyndep, utils, workspace
 from hypothesis import given
 import hypothesis.strategies as st
 
@@ -40,8 +40,8 @@ def test_convolution_gradients(self, pad, kernel, size,
             - 0.5
         b = np.random.rand(channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, w, b] if use_bias else [X, w]
         # Error handling path.

caffe2/python/operator_test/group_conv_test.py

@@ -7,7 +7,7 @@
 import hypothesis.strategies as st
 
 from caffe2.proto import caffe2_pb2
-from caffe2.python import core
+from caffe2.python import core, utils
 import caffe2.python.hip_test_util as hiputl
 import caffe2.python.hypothesis_test_util as hu
 
@@ -65,8 +65,8 @@ def test_group_convolution(
             - 0.5
         b = np.random.rand(output_channels).astype(np.float32) - 0.5
         if order == "NCHW":
-            X = X.transpose((0, 3, 1, 2))
-            w = w.transpose((0, 3, 1, 2))
+            X = utils.NHWC2NCHW(X)
+            w = utils.NHWC2NCHW(w)
 
         inputs = [X, w, b] if use_bias else [X, w]
 

caffe2/python/operator_test/instance_norm_test.py

@@ -6,7 +6,7 @@
 from hypothesis import given, assume
 import hypothesis.strategies as st
 
-from caffe2.python import core, model_helper, brew
+from caffe2.python import core, model_helper, brew, utils
 import caffe2.python.hypothesis_test_util as hu
 import caffe2.python.serialized_test.serialized_test_util as serial
 
@@ -17,14 +17,12 @@
 class TestInstanceNorm(serial.SerializedTestCase):
 
     def _get_inputs(self, N, C, H, W, order):
-        if order == 'NCHW':
-            input_data = np.random.rand(N, C, H, W).astype(np.float32)
-        elif order == 'NHWC':
+        input_data = np.random.rand(N, C, H, W).astype(np.float32)
+        if order == 'NHWC':
             # Allocate in the same order as NCHW and transpose to make sure
             # the inputs are identical on freshly-seeded calls.
-            input_data = np.random.rand(N, C, H, W).astype(np.float32)
-            input_data = np.transpose(input_data, axes=(0, 2, 3, 1))
-        else:
+            input_data = utils.NCHW2NHWC(input_data)
+        elif order != "NCHW":
             raise Exception('unknown order type ({})'.format(order))
 
         scale_data = np.random.rand(C).astype(np.float32)
@@ -128,7 +126,7 @@ def test_instance_norm_layout(self, gc, dc, N, C, H, W, store_mean,
             outputs[order] = self.ws.blobs['output'].fetch()
         np.testing.assert_allclose(
             outputs['NCHW'],
-            outputs['NHWC'].transpose((0, 3, 1, 2)),
+            utils.NHWC2NCHW(outputs["NHWC"]),
             atol=1e-4,
             rtol=1e-4)
 
@@ -166,7 +164,7 @@ def test_instance_norm_reference_check(
 
         def ref(input_blob, scale_blob, bias_blob):
             if order == 'NHWC':
-                input_blob = np.transpose(input_blob, axes=(0, 3, 1, 2))
+                input_blob = utils.NHWC2NCHW(input_blob)
 
             mean_blob = input_blob.reshape((N, C, -1)).mean(axis=2)
             inv_stdev_blob = 1.0 / \
@@ -180,8 +178,7 @@ def ref(input_blob, scale_blob, bias_blob):
                 + bias_bc
 
             if order == 'NHWC':
-                normalized_blob = np.transpose(
-                    normalized_blob, axes=(0, 2, 3, 1))
+                normalized_blob = utils.NCHW2NHWC(normalized_blob)
 
             if not store_mean and not store_inv_stdev:
                 return normalized_blob,
@@ -245,7 +242,7 @@ def test_instance_norm_model_helper(
 
         input_blob = np.random.rand(N, C, H, W).astype(np.float32)
         if order == 'NHWC':
-            input_blob = np.transpose(input_blob, axes=(0, 2, 3, 1))
+            input_blob = utils.NCHW2NHWC(input_blob)
 
         self.ws.create_blob('input').feed(input_blob)
 
@@ -262,11 +259,10 @@ def test_instance_norm_model_helper(
 
         output_blob = self.ws.blobs['output'].fetch()
         if order == 'NHWC':
-            output_blob = np.transpose(output_blob, axes=(0, 3, 1, 2))
+            output_blob = utils.NHWC2NCHW(output_blob)
 
         assert output_blob.shape == (N, C, H, W)
 
 
 if __name__ == '__main__':
-    import unittest
     unittest.main()