diff --git a/paddle/fluid/operators/generator/generate_sparse_op.py b/paddle/fluid/operators/generator/generate_sparse_op.py index 6c9f70b9283f88779ca175d89a01c87fd8296d25..10ee034ff3b47305a323850a07ec5c901c98b058 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 3e654210b9008280e12a9429c17a952991e16002..949a6c4c19b12cdaa636b41f96dc38e9db5c9e0f 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 545042b6f073e447b4c6e7e342c0be82e791e3b6..2208d34ecf4cda1328342e9f53b5cf00fd5aed08 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 03e621cc65325bf2394c255ae97b9874020751a0..41656280c500e3522aec0656d68369711c157526 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 5714514b61355c815b07e50abcda16c5692e3119..eb242a72c54b610a29ba7c660c7af03fdd4aebd2 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,