diff --git a/paddle/phi/api/lib/api_custom_impl.cc b/paddle/phi/api/lib/api_custom_impl.cc index 40f5b8b2975083808938b7612ed362610562189f..b816204c1a3996d71c57845995f6576c23b6897c 100644 --- a/paddle/phi/api/lib/api_custom_impl.cc +++ b/paddle/phi/api/lib/api_custom_impl.cc @@ -410,5 +410,153 @@ std::vector stack_grad_impl(const std::vector& x, return x_grad; } +std::vector meshgrid_impl(const std::vector& inputs) { + Backend kernel_backend = Backend::UNDEFINED; + DataLayout kernel_layout = DataLayout::UNDEFINED; + DataType kernel_data_type = DataType::UNDEFINED; + + if (kernel_backend == Backend::UNDEFINED || + kernel_layout == DataLayout::UNDEFINED || + kernel_data_type == DataType::UNDEFINED) { + auto kernel_key_set = ParseKernelKeyByInputArgs(inputs); + auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey(); + if (kernel_backend == Backend::UNDEFINED) { + kernel_backend = kernel_key.backend(); + } + if (kernel_layout == DataLayout::UNDEFINED) { + kernel_layout = kernel_key.layout(); + } + if (kernel_data_type == DataType::UNDEFINED) { + kernel_data_type = kernel_key.dtype(); + } + } + + const auto& kernel = phi::KernelFactory::Instance().SelectKernelOrThrowError( + "meshgrid", {kernel_backend, kernel_layout, kernel_data_type}); + VLOG(6) << "meshgrid API kernel key: [" << kernel_backend << ", " + << kernel_layout << ", " << kernel_data_type << "]"; + VLOG(6) << "meshgrid API kernel: " << kernel; + + auto* dev_ctx = GetDeviceContextByBackend(kernel_backend); + + auto input_inputs_vec = PrepareData(inputs, kernel.InputAt(0), {}); + std::vector input_inputs(input_inputs_vec->size()); + for (size_t i = 0; i < input_inputs.size(); ++i) { + input_inputs[i] = &input_inputs_vec->at(i); + } + + auto x_meta_vec = MakeMetaTensor(input_inputs); + std::vector inputs_metas(x_meta_vec.size()); + for (size_t i = 0; i < x_meta_vec.size(); ++i) { + inputs_metas[i] = &x_meta_vec[i]; + } + + // Calculate the number of out tensors + size_t out_number = inputs.size(); + + std::vector out; + auto dense_outs = SetKernelOutput(out_number, kernel_backend, &out); + + std::vector meta_outs; + meta_outs.reserve(out_number); + std::vector meta_out_ptrs; + meta_out_ptrs.reserve(out_number); + for (size_t i = 0; i < out_number; ++i) { + meta_outs.push_back(dense_outs[i]); + meta_out_ptrs.push_back(&meta_outs.back()); + } + phi::MeshgridInferMeta(inputs_metas, meta_out_ptrs); + + using kernel_signature = void (*)(const platform::DeviceContext&, + const std::vector&, + std::vector&); + auto* kernel_fn = kernel.GetVariadicKernelFn(); + (*kernel_fn)(*dev_ctx, input_inputs, dense_outs); + + return out; +} + +std::vector meshgrid_grad_impl( + const std::vector& inputs, + const std::vector& outputs_grad) { + Backend kernel_backend = Backend::UNDEFINED; + DataLayout kernel_layout = DataLayout::UNDEFINED; + DataType kernel_data_type = DataType::UNDEFINED; + + if (kernel_backend == Backend::UNDEFINED || + kernel_layout == DataLayout::UNDEFINED || + kernel_data_type == DataType::UNDEFINED) { + auto kernel_key_set = ParseKernelKeyByInputArgs(inputs, outputs_grad); + auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey(); + if (kernel_backend == Backend::UNDEFINED) { + kernel_backend = kernel_key.backend(); + } + if (kernel_layout == DataLayout::UNDEFINED) { + kernel_layout = kernel_key.layout(); + } + if (kernel_data_type == DataType::UNDEFINED) { + kernel_data_type = kernel_key.dtype(); + } + } + + const auto& kernel = phi::KernelFactory::Instance().SelectKernelOrThrowError( + "meshgrid_grad", {kernel_backend, kernel_layout, kernel_data_type}); + VLOG(6) << "meshgrid_grad API kernel key: [" << kernel_backend << ", " + << kernel_layout << ", " << kernel_data_type << "]"; + VLOG(6) << "meshgrid_grad API kernel: " << kernel; + + auto* dev_ctx = GetDeviceContextByBackend(kernel_backend); + + auto input_inputs_vec = PrepareData(inputs, kernel.InputAt(0), {}); + std::vector input_inputs(input_inputs_vec->size()); + for (size_t i = 0; i < input_inputs.size(); ++i) { + input_inputs[i] = &input_inputs_vec->at(i); + } + auto input_outputs_grad_vec = + PrepareData(outputs_grad, kernel.InputAt(1), {}); + std::vector input_outputs_grad( + input_outputs_grad_vec->size()); + for (size_t i = 0; i < input_outputs_grad.size(); ++i) { + input_outputs_grad[i] = &input_outputs_grad_vec->at(i); + } + + size_t out_number = inputs.size(); + std::vector api_output; + auto kernel_out = SetKernelOutput(out_number, kernel_backend, &api_output); + + auto inputs_meta_vec = MakeMetaTensor(input_inputs); + std::vector inputs_metas(inputs_meta_vec.size()); + for (size_t i = 0; i < inputs_meta_vec.size(); ++i) { + inputs_metas[i] = &inputs_meta_vec[i]; + } + + auto outputs_grad_meta_vec = MakeMetaTensor(input_outputs_grad); + std::vector outputs_grad_metas( + outputs_grad_meta_vec.size()); + for (size_t i = 0; i < outputs_grad_meta_vec.size(); ++i) { + outputs_grad_metas[i] = &outputs_grad_meta_vec[i]; + } + + std::vector meta_outs; + meta_outs.reserve(out_number); + std::vector meta_out_ptrs; + meta_out_ptrs.reserve(out_number); + for (size_t i = 0; i < out_number; ++i) { + meta_outs.push_back(kernel_out[i]); + meta_out_ptrs.push_back(&meta_outs.back()); + } + + phi::MeshgridGradInferMeta(inputs_metas, outputs_grad_metas, meta_out_ptrs); + + using kernel_signature = void (*)(const platform::DeviceContext&, + const std::vector&, + const std::vector&, + std::vector&); + auto* kernel_fn = kernel.GetVariadicKernelFn(); + (*kernel_fn)(*dev_ctx, input_inputs, input_outputs_grad, kernel_out); + + return api_output; +} + } // namespace experimental } // namespace paddle diff --git a/paddle/phi/api/lib/api_custom_impl.h b/paddle/phi/api/lib/api_custom_impl.h index 25d70d6477de1beae2db2591ae82595c318e1b4e..430eccdf430e0d2f4ceab0c1eb3669e48b4ad334 100644 --- a/paddle/phi/api/lib/api_custom_impl.h +++ b/paddle/phi/api/lib/api_custom_impl.h @@ -59,6 +59,9 @@ std::vector concat_grad_impl(const std::vector& x, std::vector stack_grad_impl(const std::vector& x, const Tensor& out_grad, int axis); +std::vector meshgrid_impl(const std::vector& inputs); +std::vector meshgrid_grad_impl(const std::vector& inputs, + const std::vector& outputs_grad); } // namespace experimental } // namespace paddle diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc index 9ee472c5c88436f4f7f68ad1e028660ef08d616c..64acc887b42c09b4d2736821f8b27172e4b15ac0 100644 --- a/paddle/phi/infermeta/backward.cc +++ b/paddle/phi/infermeta/backward.cc @@ -245,6 +245,20 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x, dx->share_meta(x); } +void MeshgridGradInferMeta(const std::vector& inputs, + const std::vector& outputs_grad, + std::vector inputs_grad) { + PADDLE_ENFORCE_GT(outputs_grad.size(), + 1, + errors::InvalidArgument( + "Number of Inputs(Out@Grad) should be larger than 1." + "But received Inputs(Out@Grad)' size = %d .", + outputs_grad.size())); + for (size_t i = 0; i < inputs.size(); i++) { + inputs_grad[i]->share_meta(*inputs[i]); + } +} + void NllLossGradInferMeta(const MetaTensor& x, const MetaTensor& label, paddle::optional weight, diff --git a/paddle/phi/infermeta/backward.h b/paddle/phi/infermeta/backward.h index fb13b4281ae6e0df9b4cf17eee997dcacc2bb60b..c0eb478168988f6fdbb782049974ebb0336cd4a8 100644 --- a/paddle/phi/infermeta/backward.h +++ b/paddle/phi/infermeta/backward.h @@ -115,6 +115,10 @@ void MaxPoolWithIndexGradInferMeta(const MetaTensor& x, bool adaptive, MetaTensor* dx); +void MeshgridGradInferMeta(const std::vector& inputs, + const std::vector& outputs_grad, + std::vector inputs_grad); + void NllLossGradInferMeta(const MetaTensor& input, const MetaTensor& label, paddle::optional weight, diff --git a/python/paddle/fluid/tests/unittests/test_meshgrid_op.py b/python/paddle/fluid/tests/unittests/test_meshgrid_op.py index 2cb83eba3767c92f0c6a43e6480e7825d1302051..95acdbe4a0687de873684b125d981aff1f78dbfc 100644 --- a/python/paddle/fluid/tests/unittests/test_meshgrid_op.py +++ b/python/paddle/fluid/tests/unittests/test_meshgrid_op.py @@ -20,6 +20,7 @@ from op_test import OpTest, skip_check_grad_ci import paddle.fluid as fluid import paddle from paddle.fluid import compiler, Program, program_guard, core +from paddle.fluid.framework import _test_eager_guard class TestMeshgridOp(OpTest): @@ -149,6 +150,10 @@ class TestMeshgridOp6(unittest.TestCase): assert np.array_equal(res_3.shape, [100, 200]) assert np.array_equal(res_4.shape, [100, 200]) + def test_api_eager_dygraph(self): + with _test_eager_guard(): + self.test_api_with_dygraph() + class TestMeshgridOp7(unittest.TestCase): def test_api_with_dygraph_list_input(self): @@ -163,6 +168,10 @@ class TestMeshgridOp7(unittest.TestCase): assert np.array_equal(res_3.shape, [100, 200]) assert np.array_equal(res_4.shape, [100, 200]) + def test_api_eager_dygraph(self): + with _test_eager_guard(): + self.test_api_with_dygraph_list_input() + class TestMeshgridOp8(unittest.TestCase): def test_api_with_dygraph_tuple_input(self): @@ -177,6 +186,40 @@ class TestMeshgridOp8(unittest.TestCase): assert np.array_equal(res_3.shape, [100, 200]) assert np.array_equal(res_4.shape, [100, 200]) + def test_api_eager_dygraph(self): + with _test_eager_guard(): + self.test_api_with_dygraph_tuple_input() + + +class TestMeshgridEager(unittest.TestCase): + def test_dygraph_final_state_api(self): + input_1 = np.random.randint(0, 100, [100, ]).astype('int32') + input_2 = np.random.randint(0, 100, [200, ]).astype('int32') + + with fluid.dygraph.guard(): + tensor_1 = fluid.dygraph.to_variable(input_1) + tensor_2 = fluid.dygraph.to_variable(input_2) + tensor_1.stop_gradient = False + tensor_2.stop_gradient = False + res_1, res_2 = paddle.tensor.meshgrid((tensor_1, tensor_2)) + sum = paddle.add_n([res_1, res_2]) + sum.backward() + with _test_eager_guard(): + tensor_eager_1 = fluid.dygraph.to_variable(input_1) + tensor_eager_2 = fluid.dygraph.to_variable(input_2) + tensor_eager_1.stop_gradient = False + tensor_eager_2.stop_gradient = False + res_eager_1, res_eager_2 = paddle.tensor.meshgrid( + (tensor_eager_1, tensor_eager_2)) + sum_eager = paddle.add_n([res_eager_1, res_eager_2]) + sum_eager.backward() + self.assertEqual(( + tensor_1.grad.numpy() == tensor_eager_1.grad.numpy()).all(), + True) + self.assertEqual(( + tensor_2.grad.numpy() == tensor_eager_2.grad.numpy()).all(), + True) + if __name__ == '__main__': paddle.enable_static() diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py index 166ae58a1977026770a2a5dd22dc42a954a58178..95f145cf447b5dfa3202cc8bb41431cd1d349433 100644 --- a/python/paddle/tensor/creation.py +++ b/python/paddle/tensor/creation.py @@ -776,10 +776,12 @@ def meshgrid(*args, **kwargs): if len(args) == 1 and isinstance(args[0], (list, tuple)): args = args[0] - if paddle.in_dynamic_mode(): + if _in_legacy_dygraph(): num = len(args) out = _C_ops.meshgrid(list(args), num) return out + if in_dygraph_mode(): + return _C_ops.final_state_meshgrid(list(args)) name = kwargs.get("name", None) helper = LayerHelper('meshgrid', **locals()) diff --git a/python/paddle/utils/code_gen/api.yaml b/python/paddle/utils/code_gen/api.yaml index a27b4115f1461782ed9d7e985883b630297e78dc..a0c484f6562c20145b099471dd2a79cb8dd45e0b 100644 --- a/python/paddle/utils/code_gen/api.yaml +++ b/python/paddle/utils/code_gen/api.yaml @@ -1120,6 +1120,12 @@ func : mean backward : mean_grad +- api : meshgrid + args : (Tensor[] inputs) + output : Tensor[] + invoke : meshgrid_impl(inputs) + backward : meshgrid_grad + - api : min args : (Tensor x, int64_t[] dims={}, bool keep_dim=false) output : Tensor(out) diff --git a/python/paddle/utils/code_gen/backward.yaml b/python/paddle/utils/code_gen/backward.yaml index 733a5052fc08bfc6466efbf349c9c53c2feb0297..5908e05a514d701a5be1c6d24e690e56567d62c8 100644 --- a/python/paddle/utils/code_gen/backward.yaml +++ b/python/paddle/utils/code_gen/backward.yaml @@ -777,6 +777,12 @@ kernel : func : mean_grad +- backward_api : meshgrid_grad + forward : meshgrid (Tensor[] inputs) -> Tensor[](outputs) + args : (Tensor[] inputs, Tensor[] outputs_grad) + output : Tensor[](inputs_grad) + invoke : meshgrid_grad_impl(inputs, outputs_grad) + - backward_api : min_grad forward: min (Tensor x, int64_t[] dims={}, bool keep_dim=false) -> Tensor(out) args : (Tensor x, Tensor out, Tensor out_grad, int64_t[] dims={}, bool keep_dim=false, bool reduce_all=false)