[OpAttr]output_size of unpool support Tensor type (#45543)

* [OpAttr]output_size of unpool support Tensor type

* fix coverage

* fix contain_var

* fix coverage

paddle/fluid/operators/unpool_op.cc

@@ -61,7 +61,8 @@ class Unpool2dOpMaker : public framework::OpProtoAndCheckerMaker {
         .InEnum({"max"});
     AddAttr<std::vector<int>>("output_size",
                               "(vector, optional). The shape of output.")
-        .SetDefault({0, 0});
+        .SetDefault({0, 0})
+        .SupportTensor();
     AddAttr<std::string>(
         "data_format",
         "(string, default NCHW) Only used in "

paddle/phi/api/yaml/legacy_api.yaml

@@ -2972,7 +2972,7 @@
   backward: uniform_random_inplace_grad
 
 - api: unpool
-  args: (Tensor x, Tensor indices, int[] ksize, int[] strides, int[] padding, int[] output_size, str data_format)
+  args: (Tensor x, Tensor indices, int[] ksize, int[] strides, int[] padding, IntArray output_size, str data_format)
   output: Tensor(out)
   infer_meta:
     func: UnpoolInferMeta

paddle/phi/api/yaml/legacy_backward.yaml

@@ -2734,8 +2734,8 @@
     data_type: x
 
 - backward_api: unpool_grad
-  forward: unpool (Tensor x, Tensor indices, int[] ksize, int[] strides, int[] padding,  int[] output_size, str data_format) -> Tensor(out)
-  args: (Tensor x, Tensor indices, Tensor out, Tensor out_grad, int[] ksize, int[] strides, int[] padding, int[] output_size, str data_format)
+  forward: unpool (Tensor x, Tensor indices, int[] ksize, int[] strides, int[] padding,  IntArray output_size, str data_format) -> Tensor(out)
+  args: (Tensor x, Tensor indices, Tensor out, Tensor out_grad, int[] ksize, int[] strides, int[] padding, IntArray output_size, str data_format)
   output: Tensor(x_grad)
   infer_meta:
     func: UnchangedInferMeta

paddle/phi/infermeta/binary.cc

@@ -2756,7 +2756,7 @@ void UnpoolInferMeta(const MetaTensor& x,
                      const std::vector<int>& ksize,
                      const std::vector<int>& strides,
                      const std::vector<int>& paddings,
-                     const std::vector<int>& output_size,
+                     const IntArray& output_size,
                      const std::string& data_format,
                      MetaTensor* out,
                      MetaConfig config) {
@@ -2780,11 +2780,16 @@ void UnpoolInferMeta(const MetaTensor& x,
                         in_y_dims));
 
   std::vector<int64_t> output_shape({in_x_dims[0], in_x_dims[1]});
+
+  std::vector<int64_t> output_size_val(output_size.size(), -1);
+  if (config.is_runtime || !output_size.FromTensor()) {
+    output_size_val = output_size.GetData();
+  }
   for (size_t i = 0; i < ksize.size(); ++i) {
     if (!config.is_runtime && in_x_dims[i + 2] <= 0) {
       output_shape.push_back(-1);
     } else {
-      output_shape.push_back(output_size[i]);
+      output_shape.push_back(output_size_val[i]);
     }
   }
   if (out != nullptr) {

paddle/phi/infermeta/binary.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 
+#include "paddle/phi/common/int_array.h"
 #include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/meta_tensor.h"
 
@@ -405,7 +406,7 @@ void UnpoolInferMeta(const MetaTensor& x,
                      const std::vector<int>& ksize,
                      const std::vector<int>& strides,
                      const std::vector<int>& paddings,
-                     const std::vector<int>& output_size,
+                     const IntArray& output_size,
                      const std::string& data_format,
                      MetaTensor* out,
                      MetaConfig config = MetaConfig());

paddle/phi/kernels/cpu/unpool_grad_kernel.cc

@@ -33,7 +33,7 @@ void UnpoolGradKernel(const Context& dev_ctx,
                       const std::vector<int>& ksize,
                       const std::vector<int>& strides,
                       const std::vector<int>& paddings,
-                      const std::vector<int>& output_size,
+                      const IntArray& output_size,
                       const std::string& data_format,
                       DenseTensor* x_grad) {
   T* input_grad_data = dev_ctx.template Alloc<T>(x_grad);

paddle/phi/kernels/cpu/unpool_kernel.cc

@@ -30,7 +30,7 @@ void UnpoolKernel(const Context& dev_ctx,
                   const std::vector<int>& ksize,
                   const std::vector<int>& strides,
                   const std::vector<int>& paddings,
-                  const std::vector<int>& output_size,
+                  const IntArray& output_size,
                   const std::string& data_format,
                   DenseTensor* out) {
   T* output_data = dev_ctx.template Alloc<T>(out);

paddle/phi/kernels/gpu/unpool_grad_kernel.cu

@@ -163,7 +163,7 @@ void UnpoolGradKernel(const Context& dev_ctx,
                       const std::vector<int>& ksize,
                       const std::vector<int>& strides,
                       const std::vector<int>& paddings,
-                      const std::vector<int>& output_size,
+                      const IntArray& output_size,
                       const std::string& data_format,
                       DenseTensor* x_grad) {
   T* input_grad_data = dev_ctx.template Alloc<T>(x_grad);

paddle/phi/kernels/gpu/unpool_kernel.cu

@@ -32,38 +32,37 @@ __global__ void KernelUnpool2dMax(const int nthreads,
                                   const int channels,
                                   T* output_data,
                                   const int output_height,
-                                  const int output_width){
-    CUDA_KERNEL_LOOP(linearIndex, nthreads){
-        int c = (linearIndex / input_width / input_height) % channels;
-int n = linearIndex / input_width / input_height / channels;
-output_data += (n * channels + c) * output_height * output_width;
-int maxind = indices_data[linearIndex];
-output_data[maxind] = input_data[linearIndex];
-}  // namespace phi
+                                  const int output_width) {
+  CUDA_KERNEL_LOOP(linearIndex, nthreads) {
+    int c = (linearIndex / input_width / input_height) % channels;
+    int n = linearIndex / input_width / input_height / channels;
+    output_data += (n * channels + c) * output_height * output_width;
+    int maxind = indices_data[linearIndex];
+    output_data[maxind] = input_data[linearIndex];
+  }
 }
-;
 
 template <typename T>
-__global__ void KernelUnpool3dMax(
-    const int nthreads,
-    const T* input_data,
-    const int* indices_data,
-    const int input_depth,
-    const int input_height,
-    const int input_width,
-    const int channels,
-    T* output_data,
-    const int output_depth,
-    const int output_height,
-    const int output_width){CUDA_KERNEL_LOOP(linearIndex, nthreads){
+__global__ void KernelUnpool3dMax(const int nthreads,
+                                  const T* input_data,
+                                  const int* indices_data,
+                                  const int input_depth,
+                                  const int input_height,
+                                  const int input_width,
+                                  const int channels,
+                                  T* output_data,
+                                  const int output_depth,
+                                  const int output_height,
+                                  const int output_width) {
+  CUDA_KERNEL_LOOP(linearIndex, nthreads) {
     int c = (linearIndex / input_depth / input_width / input_height) % channels;
-int n = linearIndex / input_depth / input_width / input_height / channels;
-output_data += (n * channels + c) * output_depth * output_height * output_width;
-int maxind = indices_data[linearIndex];
-output_data[maxind] = input_data[linearIndex];
-}
+    int n = linearIndex / input_depth / input_width / input_height / channels;
+    output_data +=
+        (n * channels + c) * output_depth * output_height * output_width;
+    int maxind = indices_data[linearIndex];
+    output_data[maxind] = input_data[linearIndex];
+  }
 }
-;
 
 template <typename T, typename Context>
 class Unpool2dMaxFunctor {
@@ -146,7 +145,7 @@ void UnpoolKernel(const Context& dev_ctx,
                   const std::vector<int>& ksize,
                   const std::vector<int>& strides,
                   const std::vector<int>& paddings,
-                  const std::vector<int>& output_size,
+                  const IntArray& output_size,
                   const std::string& data_format,
                   DenseTensor* out) {
   T* output_data = dev_ctx.template Alloc<T>(out);

paddle/phi/kernels/unpool_grad_kernel.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/phi/common/int_array.h"
 #include "paddle/phi/core/dense_tensor.h"
 
 namespace phi {
@@ -27,7 +28,7 @@ void UnpoolGradKernel(const Context& dev_ctx,
                       const std::vector<int>& ksize,
                       const std::vector<int>& strides,
                       const std::vector<int>& paddings,
-                      const std::vector<int>& output_size,
+                      const IntArray& output_size,
                       const std::string& data_format,
                       DenseTensor* x_grad);
 

paddle/phi/kernels/unpool_kernel.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/phi/common/int_array.h"
 #include "paddle/phi/core/dense_tensor.h"
 
 namespace phi {
@@ -25,7 +26,7 @@ void UnpoolKernel(const Context& dev_ctx,
                   const std::vector<int>& ksize,
                   const std::vector<int>& strides,
                   const std::vector<int>& paddings,
-                  const std::vector<int>& output_size,
+                  const IntArray& output_size,
                   const std::string& data_format,
                   DenseTensor* out);
 

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

@@ -14,10 +14,15 @@
 
 from __future__ import print_function
 
+import os
 import unittest
 import numpy as np
 from op_test import OpTest
 import paddle
+import paddle.nn.functional as F
+from paddle.fluid import Program, program_guard
+
+from test_attribute_var import UnittestBase
 
 
 def _unpool_output_size(x, kernel_size, stride, padding, output_size):
@@ -354,5 +359,83 @@ class TestUnpoolOpAPI_st(unittest.TestCase):
         np.testing.assert_allclose(results[0], expect_res, rtol=1e-05)
 
 
+class TestOutputSizeTensor(UnittestBase):
+
+    def init_info(self):
+        self.shapes = [[1, 3, 6, 6]]
+        self.save_path = os.path.join(self.temp_dir.name, self.path_prefix())
+
+    def test_static(self):
+        main_prog = Program()
+        starup_prog = Program()
+        with program_guard(main_prog, starup_prog):
+            fc = paddle.nn.Linear(6, 6)
+            x = paddle.randn(self.shapes[0])
+            x.stop_gradient = False
+            feat = fc(x)  # [1,3,6,6]
+
+            out = self.call_func(feat)
+
+            sgd = paddle.optimizer.SGD()
+            sgd.minimize(paddle.mean(out))
+            self.assertTrue(self.var_prefix() in str(main_prog))
+
+            exe = paddle.static.Executor()
+            exe.run(starup_prog)
+            res = exe.run(fetch_list=[out])
+            np.testing.assert_array_equal(res[0].shape, [1, 3, 7, 7])
+            paddle.static.save_inference_model(self.save_path, [x], [out], exe)
+            # Test for Inference Predictor
+            infer_outs = self.infer_prog()
+            np.testing.assert_array_equal(res[0].shape, [1, 3, 7, 7])
+
+    def path_prefix(self):
+        return 'unpool_var'
+
+    def var_prefix(self):
+        return "Vars["
+
+    def call_func(self, x):
+        output_size = [paddle.assign([7]), paddle.assign([7])]
+        pool_out, indices = F.max_pool2d(x,
+                                         kernel_size=2,
+                                         stride=2,
+                                         padding=0,
+                                         return_mask=True)
+        # pool_out shape: [1, 1, 6, 6],  indices shape: [1, 1, 6, 6]
+        unpool_out = F.max_unpool2d(pool_out,
+                                    indices,
+                                    kernel_size=2,
+                                    padding=0,
+                                    output_size=output_size)
+        # unpool_out shape: [1, 1, 7, 7]
+        return unpool_out
+
+
+class TestZOutputSizeTensor2(unittest.TestCase):
+
+    def setUp(self):
+        paddle.disable_static()
+
+    def tearDown(self):
+        paddle.enable_static()
+
+    def test_dygraph(self):
+        x = paddle.randn([1, 3, 6, 6])
+        pool_out, indices = F.max_pool2d(x,
+                                         kernel_size=2,
+                                         stride=2,
+                                         padding=0,
+                                         return_mask=True)
+        output_size = [paddle.assign([7]), paddle.assign([7])]
+        unpool_out = F.max_unpool2d(pool_out,
+                                    indices,
+                                    kernel_size=2,
+                                    padding=0,
+                                    output_size=output_size)
+        np.testing.assert_array_equal(unpool_out.shape, [1, 3, 7, 7])
+
+
 if __name__ == '__main__':
+    paddle.enable_static()
     unittest.main()

python/paddle/nn/functional/pooling.py

@@ -18,8 +18,9 @@ from ...tensor.manipulation import unsqueeze, squeeze
 from ...fluid.data_feeder import check_type, check_variable_and_dtype
 from paddle import _C_ops, _legacy_C_ops
 from paddle import in_dynamic_mode
-from paddle.fluid.framework import _in_legacy_dygraph
-from paddle.fluid.framework import in_dygraph_mode
+from paddle.fluid import core
+from paddle.fluid.framework import _in_legacy_dygraph, Variable
+from paddle.fluid.framework import in_dygraph_mode, _non_static_mode
 
 __all__ = []
 
@@ -651,8 +652,19 @@ def _unpool_output_size(x, kernel_size, stride, padding, output_size):
     for d in range(len(kernel_size)):
         default_size.append((input_size[-len(kernel_size) + d] - 1) *
                             stride[d] + kernel_size[d] - 2 * padding[d])
+
+    has_static_var = False
     if output_size is None:
         ret = default_size
+    elif utils._contain_var(output_size):
+        if not _non_static_mode():
+            has_static_var = True
+            output_size = utils._convert_to_tensor_list(output_size)
+        else:
+            for i, var in enumerate(output_size):
+                if isinstance(var, Variable):
+                    output_size[i] = var.numpy()[0]
+        ret = output_size
     else:
         if len(output_size) == len(kernel_size) + 2:
             output_size = output_size[2:]
@@ -662,13 +674,14 @@ def _unpool_output_size(x, kernel_size, stride, padding, output_size):
                 "{} or {} elements, but it has a length of '{}'".format(
                     len(kernel_size),
                     len(kernel_size) + 2, len(output_size)))
-        for d in range(len(kernel_size)):
-            min_size = default_size[d] - stride[d]
-            max_size = default_size[d] + stride[d]
-            if not (min_size < output_size[d] < max_size):
-                raise ValueError(
-                    'invalid output_size "{}" (dim {} must be between {} and {})'
-                    .format(output_size, d, min_size, max_size))
+        if not has_static_var:
+            for d in range(len(kernel_size)):
+                min_size = default_size[d] - stride[d]
+                max_size = default_size[d] + stride[d]
+                if not (min_size < output_size[d] < max_size):
+                    raise ValueError(
+                        'invalid output_size "{}" (dim {} must be between {} and {})'
+                        .format(output_size, d, min_size, max_size))
 
         ret = output_size
     return ret