Merge pull request #2 from qingqing01/pixel_shuffle

Fix unit testing

Author: qingqing01

paddle/fluid/operators/pixel_shuffle_op.cu

@@ -23,5 +23,5 @@ REGISTER_OP_CUDA_KERNEL(
     ops::PixelShuffleOpCUDAKernel<plat::CUDADeviceContext, double>);
 REGISTER_OP_CUDA_KERNEL(
     pixel_shuffle_grad,
-    ops::PixelShuffleGradOpCUDAKernel<plat::CUDADeviceContext, float>,
-    ops::PixelShuffleGradOpCUDAKernel<plat::CUDADeviceContext, double>);
+    ops::PixelShuffleGradOpKernel<plat::CUDADeviceContext, float>,
+    ops::PixelShuffleGradOpKernel<plat::CUDADeviceContext, double>);

paddle/fluid/operators/pixel_shuffle_op.h

@@ -29,7 +29,7 @@ class PixelShuffleOpKernel : public framework::OpKernel<T> {
     int factor = ctx.Attr<int>("upscale_factor");
 
     auto in_dims = in->dims();
-    auto o_dims = in->dims();
+    auto o_dims = out->dims();
 
     framework::Tensor t;
     t.ShareDataWith(*in);
@@ -44,6 +44,7 @@ class PixelShuffleOpKernel : public framework::OpKernel<T> {
     math::Transpose<DeviceContext, T, 6> trans;
     auto& dev_ctx = ctx.template device_context<DeviceContext>();
     trans(dev_ctx, t, &o, axis);
+    out->Resize(o_dims);
   }
 };
 
@@ -58,7 +59,7 @@ class PixelShuffleGradOpKernel : public framework::OpKernel<T> {
     int factor = ctx.Attr<int>("upscale_factor");
 
     auto do_dims = dout->dims();
-    auto dx_dims = dout->dims();
+    auto dx_dims = dx->dims();
 
     framework::Tensor t;
     t.ShareDataWith(*dout);
@@ -73,6 +74,7 @@ class PixelShuffleGradOpKernel : public framework::OpKernel<T> {
     math::Transpose<DeviceContext, T, 6> trans;
     auto& dev_ctx = ctx.template device_context<DeviceContext>();
     trans(dev_ctx, t, &o, axis);
+    dx->Resize(dx_dims);
   }
 };
 

python/paddle/fluid/tests/unittests/test_pixel_shuffle.py

@@ -25,17 +25,19 @@ def setUp(self):
         n, c, h, w = 2, 9, 4, 4
         up_factor = 3
         shape = [n, c, h, w]
-        x = np.random.random(self.shape).astype("float32")
+        x = np.random.random(shape).astype("float32")
 
         new_shape = (n, c / (up_factor * up_factor), up_factor, up_factor, h, w)
         # reshape to (num,output_channel,upscale_factor,upscale_factor,h,w)
         npresult = np.reshape(x, new_shape)
         # transpose to (num,output_channel,h,upscale_factor,w,upscale_factor)
         npresult = npresult.transpose(0, 1, 4, 2, 5, 3)
-        npresult = np.reshape(npresult, (1, 1, 12, 12))
+        oshape = [n, c / (up_factor * up_factor), h * up_factor, w * up_factor]
+        npresult = np.reshape(npresult, oshape)
 
         self.inputs = {'X': x}
         self.outputs = {'Out': npresult}
+        self.attrs = {'upscale_factor': up_factor}
 
     def test_check_output(self):
         self.check_output()