PaddlePaddle
diff --git a/‎paddle/phi/kernels/cpu/unpool_grad_kernel.cc
Copy file name to clipboard
+58-27Lines changed: 58 additions & 27 deletions b/‎paddle/phi/kernels/cpu/unpool_grad_kernel.cc
Copy file name to clipboard
+58-27Lines changed: 58 additions & 27 deletions
diff --git a/‎paddle/phi/kernels/cpu/unpool_kernel.cc
Copy file name to clipboard
+49-23Lines changed: 49 additions & 23 deletions b/‎paddle/phi/kernels/cpu/unpool_kernel.cc
Copy file name to clipboard
+49-23Lines changed: 49 additions & 23 deletions
diff --git a/‎paddle/phi/kernels/gpu/unpool_grad_kernel.cu
Copy file name to clipboard
+28-16Lines changed: 28 additions & 16 deletions b/‎paddle/phi/kernels/gpu/unpool_grad_kernel.cu
Copy file name to clipboard
+28-16Lines changed: 28 additions & 16 deletions
@@ -24,18 +24,13 @@
 
 namespace phi {
 
-template <typename T, typename Context>
-void UnpoolGradKernel(const Context& dev_ctx,
-                      const DenseTensor& x,
-                      const DenseTensor& indices,
-                      const DenseTensor& out,
-                      const DenseTensor& out_grad,
-                      const std::vector<int>& ksize UNUSED,
-                      const std::vector<int>& strides UNUSED,
-                      const std::vector<int>& paddings UNUSED,
-                      const IntArray& output_size UNUSED,
-                      const std::string& data_format UNUSED,
-                      DenseTensor* x_grad) {
+template <typename T, typename IndT, typename Context>
+void UnpoolGrad(const Context& dev_ctx,
+                const DenseTensor& x,
+                const DenseTensor& indices,
+                const DenseTensor& out,
+                const DenseTensor& out_grad,
+                DenseTensor* x_grad) {
   T* input_grad_data = dev_ctx.template Alloc<T>(x_grad);
   const T* output_grad_data = out_grad.data<T>();
   phi::funcs::SetConstant<Context, T> zero;
@@ -48,12 +43,12 @@ void UnpoolGradKernel(const Context& dev_ctx,
   const int output_width = static_cast<int>(out.dims()[3]);
   int input_feasize = input_height * input_width;
   int output_feasize = output_height * output_width;
-  const int* indices_data = indices.data<int>();
+  const IndT* indices_data = indices.data<IndT>();
 
   for (int b = 0; b < batch_size; ++b) {
     for (int c = 0; c < output_channels; ++c) {
       for (int i = 0; i < input_feasize; ++i) {
-        int index = indices_data[i];
+        IndT index = indices_data[i];
         PADDLE_ENFORCE_LT(
             index,
             output_feasize,
@@ -75,17 +70,32 @@ void UnpoolGradKernel(const Context& dev_ctx,
 }
 
 template <typename T, typename Context>
-void Unpool3dGradKernel(const Context& dev_ctx,
-                        const DenseTensor& x,
-                        const DenseTensor& indices,
-                        const DenseTensor& out,
-                        const DenseTensor& out_grad,
-                        const std::vector<int>& ksize UNUSED,
-                        const std::vector<int>& strides,
-                        const std::vector<int>& paddings UNUSED,
-                        const std::vector<int>& output_size UNUSED,
-                        const std::string& data_format UNUSED,
-                        DenseTensor* x_grad) {
+void UnpoolGradKernel(const Context& dev_ctx,
+                      const DenseTensor& x,
+                      const DenseTensor& indices,
+                      const DenseTensor& out,
+                      const DenseTensor& out_grad,
+                      const std::vector<int>& ksize UNUSED,
+                      const std::vector<int>& strides UNUSED,
+                      const std::vector<int>& paddings UNUSED,
+                      const IntArray& output_size UNUSED,
+                      const std::string& data_format UNUSED,
+                      DenseTensor* x_grad) {
+  const auto& indices_type = indices.dtype();
+  if (indices_type == phi::DataType::INT32) {
+    UnpoolGrad<T, int, Context>(dev_ctx, x, indices, out, out_grad, x_grad);
+  } else {
+    UnpoolGrad<T, int64_t, Context>(dev_ctx, x, indices, out, out_grad, x_grad);
+  }
+}
+
+template <typename T, typename IndT, typename Context>
+void Unpool3dGrad(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  const DenseTensor& indices,
+                  const DenseTensor& out,
+                  const DenseTensor& out_grad,
+                  DenseTensor* x_grad) {
   T* input_grad_data = dev_ctx.template Alloc<T>(x_grad);
   const T* output_grad_data = out_grad.data<T>();
   phi::funcs::SetConstant<Context, T> zero;
@@ -101,12 +111,12 @@ void Unpool3dGradKernel(const Context& dev_ctx,
   const int output_width = static_cast<int>(out.dims()[4]);
   int input_feasize = input_depth * input_height * input_width;
   int output_feasize = output_depth * output_height * output_width;
-  const int* indices_data = indices.data<int>();
+  const IndT* indices_data = indices.data<IndT>();
 
   for (int b = 0; b < batch_size; ++b) {
     for (int c = 0; c < output_channels; ++c) {
       for (int i = 0; i < input_feasize; ++i) {
-        int index = indices_data[i];
+        IndT index = indices_data[i];
         PADDLE_ENFORCE_LT(
             index,
             output_feasize,
@@ -128,6 +138,27 @@ void Unpool3dGradKernel(const Context& dev_ctx,
   }
 }
 
+template <typename T, typename Context>
+void Unpool3dGradKernel(const Context& dev_ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& indices,
+                        const DenseTensor& out,
+                        const DenseTensor& out_grad,
+                        const std::vector<int>& ksize UNUSED,
+                        const std::vector<int>& strides,
+                        const std::vector<int>& paddings UNUSED,
+                        const std::vector<int>& output_size UNUSED,
+                        const std::string& data_format UNUSED,
+                        DenseTensor* x_grad) {
+  const auto& indices_type = indices.dtype();
+  if (indices_type == phi::DataType::INT32) {
+    Unpool3dGrad<T, int, Context>(dev_ctx, x, indices, out, out_grad, x_grad);
+  } else {
+    Unpool3dGrad<T, int64_t, Context>(
+        dev_ctx, x, indices, out, out_grad, x_grad);
+  }
+}
+
 }  // namespace phi
 
 PD_REGISTER_KERNEL(unpool_grad,
 
@@ -23,16 +23,11 @@
 
 namespace phi {
 
-template <typename T, typename Context>
-void UnpoolKernel(const Context& dev_ctx,
-                  const DenseTensor& x,
-                  const DenseTensor& indices,
-                  const std::vector<int>& ksize UNUSED,
-                  const std::vector<int>& strides UNUSED,
-                  const std::vector<int>& paddings UNUSED,
-                  const IntArray& output_size UNUSED,
-                  const std::string& data_format UNUSED,
-                  DenseTensor* out) {
+template <typename T, typename IndT, typename Context>
+void Unpool(const Context& dev_ctx,
+            const DenseTensor& x,
+            const DenseTensor& indices,
+            DenseTensor* out) {
   T* output_data = dev_ctx.template Alloc<T>(out);
   if (output_data) {
     phi::funcs::SetConstant<Context, T> set_zero;
@@ -47,11 +42,11 @@ void UnpoolKernel(const Context& dev_ctx,
   int input_feasize = input_height * input_width;
   int output_feasize = output_height * output_width;
   const T* input_data = x.data<T>();
-  const int* indices_data = indices.data<int>();
+  const IndT* indices_data = indices.data<IndT>();
   for (int b = 0; b < batch_size; ++b) {
     for (int c = 0; c < output_channels; ++c) {
       for (int i = 0; i < input_feasize; ++i) {
-        int index = indices_data[i];
+        IndT index = indices_data[i];
         PADDLE_ENFORCE_LT(
             index,
             output_feasize,
@@ -73,15 +68,28 @@ void UnpoolKernel(const Context& dev_ctx,
 }
 
 template <typename T, typename Context>
-void Unpool3dKernel(const Context& dev_ctx,
-                    const DenseTensor& x,
-                    const DenseTensor& indices,
-                    const std::vector<int>& ksize UNUSED,
-                    const std::vector<int>& strides UNUSED,
-                    const std::vector<int>& paddings UNUSED,
-                    const std::vector<int>& output_size UNUSED,
-                    const std::string& data_format UNUSED,
-                    DenseTensor* out) {
+void UnpoolKernel(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  const DenseTensor& indices,
+                  const std::vector<int>& ksize UNUSED,
+                  const std::vector<int>& strides UNUSED,
+                  const std::vector<int>& paddings UNUSED,
+                  const IntArray& output_size UNUSED,
+                  const std::string& data_format UNUSED,
+                  DenseTensor* out) {
+  const auto& indices_type = indices.dtype();
+  if (indices_type == phi::DataType::INT32) {
+    Unpool<T, int, Context>(dev_ctx, x, indices, out);
+  } else {
+    Unpool<T, int64_t, Context>(dev_ctx, x, indices, out);
+  }
+}
+
+template <typename T, typename IndT, typename Context>
+void Unpool3d(const Context& dev_ctx,
+              const DenseTensor& x,
+              const DenseTensor& indices,
+              DenseTensor* out) {
   T* output_data = dev_ctx.template Alloc<T>(out);
   if (output_data) {
     phi::funcs::SetConstant<Context, T> set_zero;
@@ -98,11 +106,11 @@ void Unpool3dKernel(const Context& dev_ctx,
   int input_feasize = input_depth * input_height * input_width;
   int output_feasize = output_depth * output_height * output_width;
   const T* input_data = x.data<T>();
-  const int* indices_data = indices.data<int>();
+  const IndT* indices_data = indices.data<IndT>();
   for (int b = 0; b < batch_size; ++b) {
     for (int c = 0; c < output_channels; ++c) {
       for (int i = 0; i < input_feasize; ++i) {
-        int index = indices_data[i];
+        IndT index = indices_data[i];
         PADDLE_ENFORCE_LT(
             index,
             output_feasize,
@@ -124,6 +132,24 @@ void Unpool3dKernel(const Context& dev_ctx,
   }
 }
 
+template <typename T, typename Context>
+void Unpool3dKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const DenseTensor& indices,
+                    const std::vector<int>& ksize UNUSED,
+                    const std::vector<int>& strides UNUSED,
+                    const std::vector<int>& paddings UNUSED,
+                    const std::vector<int>& output_size UNUSED,
+                    const std::string& data_format UNUSED,
+                    DenseTensor* out) {
+  const auto& indices_type = indices.dtype();
+  if (indices_type == phi::DataType::INT32) {
+    Unpool3d<T, int, Context>(dev_ctx, x, indices, out);
+  } else {
+    Unpool3d<T, int64_t, Context>(dev_ctx, x, indices, out);
+  }
+}
+
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
 
@@ -23,10 +23,10 @@
 
 namespace phi {
 
-template <typename T>
+template <typename T, typename IndT>
 __global__ void KernelUnpool2dMaxGrad(const int nthreads,
                                       const T* input_data,
-                                      const int* indices_data,
+                                      const IndT* indices_data,
                                       const int input_height,
                                       const int input_width,
                                       const int channels,
@@ -39,15 +39,15 @@ __global__ void KernelUnpool2dMaxGrad(const int nthreads,
     int c = (linearIndex / input_width / input_height) % channels;
     int n = linearIndex / input_width / input_height / channels;
     output_grad += (n * channels + c) * output_height * output_width;
-    int maxind = indices_data[linearIndex];
+    IndT maxind = indices_data[linearIndex];
     input_grad[linearIndex] = output_grad[maxind];
   }
 }
 
-template <typename T>
+template <typename T, typename IndT>
 __global__ void KernelUnpool3dMaxGrad(const int nthreads,
                                       const T* input_data,
-                                      const int* indices_data,
+                                      const IndT* indices_data,
                                       const int input_depth,
                                       const int input_height,
                                       const int input_width,
@@ -63,12 +63,12 @@ __global__ void KernelUnpool3dMaxGrad(const int nthreads,
     int n = linearIndex / input_depth / input_width / input_height / channels;
     output_grad +=
         (n * channels + c) * output_depth * output_height * output_width;
-    int maxind = indices_data[linearIndex];
+    IndT maxind = indices_data[linearIndex];
     input_grad[linearIndex] = output_grad[maxind];
   }
 }
 
-template <typename T, typename Context>
+template <typename T, typename IndT, typename Context>
 class Unpool2dMaxGradFunctor {
  public:
   void operator()(const Context& dev_ctx,
@@ -84,13 +84,13 @@ class Unpool2dMaxGradFunctor {
     const int output_height = output.dims()[2];
     const int output_width = output.dims()[3];
     const T* input_data = input.data<T>();
-    const int* indices_data = indices.data<int>();
+    const IndT* indices_data = indices.data<IndT>();
     const T* output_data = output.data<T>();
     const T* output_grad_data = output_grad.data<T>();
     T* input_grad_data = dev_ctx.template Alloc<T>(input_grad);
     int threads = 1024;
     int grid = (input.numel() + threads - 1) / threads;
-    KernelUnpool2dMaxGrad<T>
+    KernelUnpool2dMaxGrad<T, IndT>
         <<<grid, threads, 0, dev_ctx.stream()>>>(input.numel(),
                                                  input_data,
                                                  indices_data,
@@ -105,7 +105,7 @@ class Unpool2dMaxGradFunctor {
   }
 };
 
-template <typename T, typename Context>
+template <typename T, typename IndT, typename Context>
 class Unpool3dMaxGradFunctor {
  public:
   void operator()(const Context& dev_ctx,
@@ -123,13 +123,13 @@ class Unpool3dMaxGradFunctor {
     const int output_height = output.dims()[3];
     const int output_width = output.dims()[4];
     const T* input_data = input.data<T>();
-    const int* indices_data = indices.data<int>();
+    const IndT* indices_data = indices.data<IndT>();
     const T* output_data = output.data<T>();
     const T* output_grad_data = output_grad.data<T>();
     T* input_grad_data = dev_ctx.template Alloc<T>(input_grad);
     int threads = 1024;
     int grid = (input.numel() + threads - 1) / threads;
-    KernelUnpool3dMaxGrad<T>
+    KernelUnpool3dMaxGrad<T, IndT>
         <<<grid, threads, 0, dev_ctx.stream()>>>(input.numel(),
                                                  input_data,
                                                  indices_data,
@@ -162,8 +162,14 @@ void UnpoolGradKernel(const Context& dev_ctx,
   const T* output_grad_data = out_grad.data<T>();
   phi::funcs::SetConstant<Context, T> zero;
   zero(dev_ctx, x_grad, static_cast<T>(0));
-  Unpool2dMaxGradFunctor<T, Context> unpool2d_max_backward;
-  unpool2d_max_backward(dev_ctx, x, indices, out, out_grad, x_grad);
+  const auto& indices_type = indices.dtype();
+  if (indices_type == phi::DataType::INT32) {
+    Unpool2dMaxGradFunctor<T, int, Context> unpool2d_max_backward;
+    unpool2d_max_backward(dev_ctx, x, indices, out, out_grad, x_grad);
+  } else {
+    Unpool2dMaxGradFunctor<T, int64_t, Context> unpool2d_max_backward;
+    unpool2d_max_backward(dev_ctx, x, indices, out, out_grad, x_grad);
+  }
 }
 
 template <typename T, typename Context>
@@ -182,8 +188,14 @@ void Unpool3dGradKernel(const Context& dev_ctx,
   const T* output_grad_data = out_grad.data<T>();
   phi::funcs::SetConstant<Context, T> zero;
   zero(dev_ctx, x_grad, static_cast<T>(0));
-  Unpool3dMaxGradFunctor<T, Context> unpool3d_max_backward;
-  unpool3d_max_backward(dev_ctx, x, indices, out, out_grad, x_grad);
+  const auto& indices_type = indices.dtype();
+  if (indices_type == phi::DataType::INT32) {
+    Unpool3dMaxGradFunctor<T, int, Context> unpool3d_max_backward;
+    unpool3d_max_backward(dev_ctx, x, indices, out, out_grad, x_grad);
+  } else {
+    Unpool3dMaxGradFunctor<T, int64_t, Context> unpool3d_max_backward;
+    unpool3d_max_backward(dev_ctx, x, indices, out, out_grad, x_grad);
+  }
 }
 
 }  // namespace phi