From c857841ef31e0907d9f5df88d0bcf3a192dfa7ef Mon Sep 17 00:00:00 2001 From: WangZhen <23097963+0x45f@users.noreply.github.com> Date: Tue, 30 Aug 2022 10:32:01 +0800 Subject: [PATCH] Adapt tensor num_samples for multinomial (#45522) --- paddle/fluid/operators/multinomial_op.cc | 10 ++++- paddle/phi/api/yaml/legacy_api.yaml | 2 +- paddle/phi/infermeta/unary.cc | 22 ++++++---- paddle/phi/infermeta/unary.h | 5 ++- paddle/phi/kernels/cpu/multinomial_kernel.cc | 4 +- paddle/phi/kernels/gpu/multinomial_kernel.cu | 13 +++--- paddle/phi/kernels/multinomial_kernel.h | 3 +- .../tests/unittests/test_multinomial_op.py | 44 +++++++++++++++++++ 8 files changed, 82 insertions(+), 21 deletions(-) diff --git a/paddle/fluid/operators/multinomial_op.cc b/paddle/fluid/operators/multinomial_op.cc index cc08c2d54f..663e861373 100644 --- a/paddle/fluid/operators/multinomial_op.cc +++ b/paddle/fluid/operators/multinomial_op.cc @@ -31,7 +31,8 @@ class MultinomialOpMaker : public framework::OpProtoAndCheckerMaker { AddInput("X", "A tensor contains probabilities of categories"); AddOutput("Out", "The output tensor of multinomial op"); AddAttr("num_samples", "number of the generated samples") - .SetDefault(1); + .SetDefault(1) + .SupportTensor(); AddAttr("replacement", "can a category be sampled more than once") .SetDefault(false); AddComment(R"DOC( @@ -46,6 +47,13 @@ This OP returns a Tensor filled with the sampled categoris according to Multinom class MultinomialOp : public framework::OperatorWithKernel { public: using framework::OperatorWithKernel::OperatorWithKernel; + + framework::OpKernelType GetExpectedKernelType( + const framework::ExecutionContext& ctx) const override { + auto input_data_type = + framework::OperatorWithKernel::IndicateVarDataType(ctx, "X"); + return framework::OpKernelType(input_data_type, ctx.GetPlace()); + } }; } // namespace operators diff --git a/paddle/phi/api/yaml/legacy_api.yaml b/paddle/phi/api/yaml/legacy_api.yaml index 6e87936e72..8f9257d124 100755 --- a/paddle/phi/api/yaml/legacy_api.yaml +++ b/paddle/phi/api/yaml/legacy_api.yaml @@ -1828,7 +1828,7 @@ optional : rois_num - api : multinomial - args : (Tensor x, int num_samples, bool replacement) + args : (Tensor x, Scalar num_samples, bool replacement) output : Tensor(out) infer_meta : func : MultinomialInferMeta diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc index 9925a10e6d..bab0b84f47 100644 --- a/paddle/phi/infermeta/unary.cc +++ b/paddle/phi/infermeta/unary.cc @@ -1905,9 +1905,11 @@ void ModeInferMeta(const MetaTensor& x, } void MultinomialInferMeta(const MetaTensor& x, - int num_samples, + const Scalar& num_samples, bool replacement, - MetaTensor* out) { + MetaTensor* out, + MetaConfig config) { + auto int_num_samples = num_samples.to(); auto x_dim = x.dims(); int64_t x_rank = x_dim.size(); PADDLE_ENFORCE_GT(x_rank, @@ -1928,12 +1930,16 @@ void MultinomialInferMeta(const MetaTensor& x, out_dims[i] = x_dim[i]; } - PADDLE_ENFORCE_GT( - num_samples, - 0, - errors::InvalidArgument( - "The number of samples should be > 0, but got %d.", num_samples)); - out_dims[x_rank - 1] = num_samples; + if (config.is_runtime || !num_samples.FromTensor()) { + PADDLE_ENFORCE_GT(int_num_samples, + 0, + errors::InvalidArgument( + "The number of samples should be > 0, but got %d.", + int_num_samples)); + out_dims[x_rank - 1] = int_num_samples; + } else { + out_dims[x_rank - 1] = -1; + } out->set_dims(make_ddim(out_dims)); out->set_dtype(DataType::INT64); diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h index c0e20714a7..db2d51b5fe 100644 --- a/paddle/phi/infermeta/unary.h +++ b/paddle/phi/infermeta/unary.h @@ -277,9 +277,10 @@ void ModeInferMeta(const MetaTensor& x, MetaTensor* indices); void MultinomialInferMeta(const MetaTensor& x, - int num_samples, + const Scalar& num_samples, bool replacement, - MetaTensor* out); + MetaTensor* out, + MetaConfig config = MetaConfig()); void NanmedianInferMeta(const MetaTensor& x, const IntArray& axes, diff --git a/paddle/phi/kernels/cpu/multinomial_kernel.cc b/paddle/phi/kernels/cpu/multinomial_kernel.cc index e9c2a569e0..1e31b4649f 100644 --- a/paddle/phi/kernels/cpu/multinomial_kernel.cc +++ b/paddle/phi/kernels/cpu/multinomial_kernel.cc @@ -23,7 +23,7 @@ namespace phi { template void MultinomialKernel(const Context& dev_ctx, const DenseTensor& x, - int num_samples, + const Scalar& num_samples, bool replacement, DenseTensor* out) { auto* in_data = x.data(); @@ -36,7 +36,7 @@ void MultinomialKernel(const Context& dev_ctx, funcs::MultinomialFunctor(dev_ctx, out_data, in_data, - num_samples, + num_samples.to(), replacement, num_categories, num_distributions); diff --git a/paddle/phi/kernels/gpu/multinomial_kernel.cu b/paddle/phi/kernels/gpu/multinomial_kernel.cu index a4fba88d2e..2e6e95f4c2 100644 --- a/paddle/phi/kernels/gpu/multinomial_kernel.cu +++ b/paddle/phi/kernels/gpu/multinomial_kernel.cu @@ -128,9 +128,10 @@ __global__ void sampleMultinomialWithReplacement( template void MultinomialKernel(const Context& dev_ctx, const DenseTensor& x, - int num_samples, + const Scalar& num_samples, bool replacement, DenseTensor* out) { + auto int_num_samples = num_samples.to(); auto* in_data = x.data(); int64_t* out_data = dev_ctx.template Alloc(out); auto in_dims = x.dims(); @@ -172,7 +173,7 @@ void MultinomialKernel(const Context& dev_ctx, } int valid_samples = num_categories - zero_num; PADDLE_ENFORCE_LE( - num_samples, + int_num_samples, valid_samples, errors::InvalidArgument("When replacement=False, 'num_samples' " "must less than or eaqual to the number of " @@ -191,14 +192,14 @@ void MultinomialKernel(const Context& dev_ctx, rand_data[idx] = in_data[idx] / rand_data[idx]; }); - if (num_samples == 1) { + if (int_num_samples == 1) { ArgMaxKernel( dev_ctx, rand, -1, true, false, 3 /*proto::VarType::INT64*/, out); } else { std::vector out_dim_vec = vectorize(out->dims()); DenseTensor value = Empty(dev_ctx, IntArray(out_dim_vec)); TopkKernel( - dev_ctx, rand, Scalar(num_samples), -1, true, true, &value, out); + dev_ctx, rand, num_samples, -1, true, true, &value, out); } return; } @@ -268,7 +269,7 @@ void MultinomialKernel(const Context& dev_ctx, int64_t device_id = dev_ctx.GetPlace().GetDeviceId(); const auto& prop = phi::backends::gpu::GetDeviceProperties(device_id); int grid_y = std::min(num_distributions, prop.maxGridSize[1]); - dim3 grid((num_samples - 1) / block.x + 1, grid_y); + dim3 grid((int_num_samples - 1) / block.x + 1, grid_y); auto gen_cuda = dev_ctx.GetGenerator(); size_t curand4_loop_times = @@ -278,7 +279,7 @@ void MultinomialKernel(const Context& dev_ctx, auto seed_offset = gen_cuda->IncrementOffset(increment); sampleMultinomialWithReplacement - <<>>(num_samples, + <<>>(int_num_samples, out_data, num_distributions, num_categories, diff --git a/paddle/phi/kernels/multinomial_kernel.h b/paddle/phi/kernels/multinomial_kernel.h index f3d8770bc1..bef22d1966 100644 --- a/paddle/phi/kernels/multinomial_kernel.h +++ b/paddle/phi/kernels/multinomial_kernel.h @@ -14,6 +14,7 @@ #pragma once +#include "paddle/phi/common/scalar.h" #include "paddle/phi/core/dense_tensor.h" namespace phi { @@ -21,7 +22,7 @@ namespace phi { template void MultinomialKernel(const Context& dev_ctx, const DenseTensor& x, - int num_samples, + const Scalar& num_samples, bool replacement, DenseTensor* out); diff --git a/python/paddle/fluid/tests/unittests/test_multinomial_op.py b/python/paddle/fluid/tests/unittests/test_multinomial_op.py index 163f9a1a07..b5e4877e05 100644 --- a/python/paddle/fluid/tests/unittests/test_multinomial_op.py +++ b/python/paddle/fluid/tests/unittests/test_multinomial_op.py @@ -21,6 +21,8 @@ from paddle.fluid import core from op_test import OpTest import numpy as np import os +from paddle.fluid import Program, program_guard +from test_attribute_var import UnittestBase def sample_output_one_dimension(out, dim): @@ -294,5 +296,47 @@ class TestRandomValue(unittest.TestCase): paddle.enable_static() +class TestMultinomialTensorNumSamples(UnittestBase): + + def init_info(self): + self.shapes = [[3, 4]] + self.save_path = os.path.join(self.temp_dir.name, self.path_prefix()) + + def path_prefix(self): + return 'multinomial_tensor_num' + + def var_prefix(self): + return "Var[" + + def call_func(self, x): + num_samples = paddle.assign(3) + out = paddle.multinomial(x, num_samples) + return out + + def test_static(self): + main_prog = Program() + starup_prog = Program() + with program_guard(main_prog, starup_prog): + fc = paddle.nn.Linear(4, 10) + x = paddle.randn([3, 4]) + x.stop_gradient = False + feat = fc(x) + out = self.call_func(paddle.abs(feat)) + sgd = paddle.optimizer.SGD() + sgd.minimize(paddle.mean(paddle.cast(out, 'float32'))) + self.assertTrue(self.var_prefix() in str(main_prog)) + + exe = paddle.static.Executor() + exe.run(starup_prog) + res = exe.run(fetch_list=[feat, out]) + paddle.static.save_inference_model(self.save_path, [x], [feat, out], + exe) + np.testing.assert_equal(res[1].shape, (3, 3)) + + # Test for Inference Predictor + infer_outs = self.infer_prog() + np.testing.assert_equal(infer_outs[1].shape, (3, 3)) + + if __name__ == "__main__": unittest.main() -- GitLab