From d33d6db0d3bea87551bce612c00c0f37005afb4e Mon Sep 17 00:00:00 2001 From: zhangkaihuo Date: Tue, 29 Nov 2022 10:13:51 +0800 Subject: [PATCH] [Sparse]BatchNorm use inplace (#48254) --- .../operators/generator/generate_sparse_op.py | 2 ++ paddle/phi/api/yaml/sparse_backward.yaml | 2 +- paddle/phi/api/yaml/sparse_ops.yaml | 4 +-- paddle/phi/kernels/sparse/batch_norm_kernel.h | 36 +++++++++---------- python/paddle/sparse/nn/layer/norm.py | 2 +- 5 files changed, 24 insertions(+), 22 deletions(-) diff --git a/paddle/fluid/operators/generator/generate_sparse_op.py b/paddle/fluid/operators/generator/generate_sparse_op.py index 6c9f70b928..10ee034ff3 100644 --- a/paddle/fluid/operators/generator/generate_sparse_op.py +++ b/paddle/fluid/operators/generator/generate_sparse_op.py @@ -82,6 +82,8 @@ def main(op_yaml_path, backward_yaml_path, output_op_path, output_arg_map_path): backward_op_dict = to_named_dict(backward_ops) for op in ops: + if op['name'][-1] == '_': + op['name'] = op['name'][:-1] op['op_name'] = SPARSE_OP_PREFIX + op['name'] op['name'] = op['op_name'] if op["backward"] is not None: diff --git a/paddle/phi/api/yaml/sparse_backward.yaml b/paddle/phi/api/yaml/sparse_backward.yaml index 3e654210b9..949a6c4c19 100644 --- a/paddle/phi/api/yaml/sparse_backward.yaml +++ b/paddle/phi/api/yaml/sparse_backward.yaml @@ -101,7 +101,7 @@ atanh_csr_grad {sparse_csr, sparse_csr -> sparse_csr} - backward_op : batch_norm_grad - forward : batch_norm (Tensor x, Tensor mean, Tensor variance, Tensor scale, Tensor bias, bool is_test, float momentum, float epsilon, str data_layout, bool use_global_stats, bool trainable_statistics) -> Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space) + forward : batch_norm_ (Tensor x, Tensor mean, Tensor variance, Tensor scale, Tensor bias, bool is_test, float momentum, float epsilon, str data_layout, bool use_global_stats, bool trainable_statistics) -> Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space) args : (Tensor x, Tensor scale, Tensor bias, Tensor mean_out, Tensor variance_out, Tensor saved_mean, Tensor saved_variance, Tensor reserve_space, Tensor out_grad, float momentum, float epsilon, str data_layout, bool is_test, bool use_global_stats, bool trainable_statistics) output : Tensor(x_grad), Tensor(scale_grad), Tensor(bias_grad) infer_meta : diff --git a/paddle/phi/api/yaml/sparse_ops.yaml b/paddle/phi/api/yaml/sparse_ops.yaml index 545042b6f0..2208d34ecf 100644 --- a/paddle/phi/api/yaml/sparse_ops.yaml +++ b/paddle/phi/api/yaml/sparse_ops.yaml @@ -87,7 +87,7 @@ layout : x backward : atanh_grad -- op : batch_norm +- op : batch_norm_ args : (Tensor x, Tensor mean, Tensor variance, Tensor scale, Tensor bias, bool is_test, float momentum, float epsilon, str data_layout, bool use_global_stats, bool trainable_statistics) output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space) infer_meta : @@ -95,7 +95,7 @@ kernel : func : batch_norm_coo {sparse_coo, dense, dense, dense, dense -> sparse_coo, dense, dense, dense, dense, dense} data_type : x - view : (mean -> mean_out), (variance -> variance_out) + inplace : (mean -> mean_out), (variance -> variance_out) backward : batch_norm_grad - op : cast diff --git a/paddle/phi/kernels/sparse/batch_norm_kernel.h b/paddle/phi/kernels/sparse/batch_norm_kernel.h index 03e621cc65..41656280c5 100644 --- a/paddle/phi/kernels/sparse/batch_norm_kernel.h +++ b/paddle/phi/kernels/sparse/batch_norm_kernel.h @@ -23,24 +23,24 @@ namespace phi { namespace sparse { template -void BatchNormKernel(const Context& dev_ctx, - const SparseCooTensor& x, - const DenseTensor& scale, - const DenseTensor& bias, - const DenseTensor& mean, - const DenseTensor& variance, - float momentum, - float epsilon, - const std::string& data_layout, - bool is_test, - bool use_global_stats, - bool trainable_statistics, - SparseCooTensor* y, - DenseTensor* mean_out, - DenseTensor* variance_out, - DenseTensor* saved_mean, - DenseTensor* saved_variance, - DenseTensor* reserve_space); +void BatchNormCooKernel(const Context& dev_ctx, + const SparseCooTensor& x, + const DenseTensor& mean, + const DenseTensor& variance, + const DenseTensor& scale, + const DenseTensor& bias, + bool is_test, + float momentum, + float epsilon, + const std::string& data_layout, + bool use_global_stats, + bool trainable_statistics, + SparseCooTensor* y, + DenseTensor* mean_out, + DenseTensor* variance_out, + DenseTensor* saved_mean, + DenseTensor* saved_variance, + DenseTensor* reserve_space); } // namespace sparse } // namespace phi diff --git a/python/paddle/sparse/nn/layer/norm.py b/python/paddle/sparse/nn/layer/norm.py index 5714514b61..eb242a72c5 100644 --- a/python/paddle/sparse/nn/layer/norm.py +++ b/python/paddle/sparse/nn/layer/norm.py @@ -138,7 +138,7 @@ class BatchNorm(paddle.nn.BatchNorm1D): data_format = 'NCHW' if self._data_format[1] == 'C' else 'NHWC' if in_dynamic_mode(): - batch_norm_out, _, _, _, _, _ = _C_ops.sparse_batch_norm( + batch_norm_out, _, _, _, _, _ = _C_ops.sparse_batch_norm_( input, self._mean, self._variance, -- GitLab