【Hackathon No.41】为 Paddle 新增 Rprop API -part

paddle/phi/api/yaml/ops.yaml

@@ -2197,6 +2197,19 @@
   inplace : (x -> out)
   backward : round_grad
 
+- op : rprop_
+  args : (Tensor param, Tensor grad, Tensor prev, Tensor learning_rate, Tensor master_param, Tensor learning_rate_range, Tensor etas, bool multi_precision=false)
+  output : Tensor(param_out), Tensor(prev_out), Tensor(learning_rate_out), Tensor(master_param_out)
+  infer_meta :
+    func : RpropInferMeta
+  kernel :
+    func : rprop
+    data_type : param
+  data_transform :
+    support_trans_dtype : learning_rate
+  optional : master_param, master_param_out
+  inplace : (param -> param_out), (prev -> prev_out), (learning_rate -> learning_rate_out), (master_param -> master_param_out)
+
 - op : rsqrt
   args : (Tensor x)
   output : Tensor(out)

paddle/phi/infermeta/multiary.cc

@@ -3470,6 +3470,50 @@ void RnnInferMeta(const MetaTensor& x,
   }
 }
 
+void RpropInferMeta(const MetaTensor& param,
+                    const MetaTensor& grad,
+                    const MetaTensor& prev,
+                    const MetaTensor& learning_rate,
+                    const MetaTensor& master_param,
+                    const MetaTensor& learning_rate_range,
+                    const MetaTensor& etas,
+                    bool multi_precision,
+                    MetaTensor* param_out,
+                    MetaTensor* prev_out,
+                    MetaTensor* learning_rate_out,
+                    MetaTensor* master_param_out) {
+  PADDLE_ENFORCE_NOT_NULL(
+      param_out,
+      phi::errors::InvalidArgument(
+          "Output(ParamOut) of RpropOp should not be null."));
+
+  PADDLE_ENFORCE_NOT_NULL(
+      prev_out,
+      phi::errors::InvalidArgument(
+          "Output(PrevOut) of RpropOp should not be null."));
+
+  PADDLE_ENFORCE_NOT_NULL(
+      learning_rate_out,
+      phi::errors::InvalidArgument(
+          "Output(LearningRateOut) of RpropOp should not be null."));
+
+  param_out->set_dims(param.dims());
+  param_out->set_dtype(param.dtype());
+  prev_out->set_dims(prev.dims());
+  prev_out->set_dtype(prev.dtype());
+  learning_rate_out->set_dims(learning_rate.dims());
+  learning_rate_out->set_dtype(learning_rate.dtype());
+  if (multi_precision) {
+    master_param_out->set_dims(master_param.dims());
+    if (DataType::FLOAT16 == master_param.dtype() ||
+        DataType::BFLOAT16 == master_param.dtype()) {
+      master_param_out->set_dtype(DataType::FLOAT32);
+    } else {
+      master_param_out->set_dtype(master_param.dtype());
+    }
+  }
+}
+
 void SgdInferMeta(const MetaTensor& param,
                   const MetaTensor& learning_rate,
                   const MetaTensor& grad,

paddle/phi/infermeta/multiary.h

@@ -627,6 +627,19 @@ void RnnInferMeta(const MetaTensor& x,
                   std::vector<MetaTensor*> state,
                   MetaTensor* reserve);
 
+void RpropInferMeta(const MetaTensor& param,
+                    const MetaTensor& grad,
+                    const MetaTensor& prev,
+                    const MetaTensor& learning_rate,
+                    const MetaTensor& master_param,
+                    const MetaTensor& learning_rate_range,
+                    const MetaTensor& etas,
+                    bool multi_precision,
+                    MetaTensor* param_out,
+                    MetaTensor* prev_out,
+                    MetaTensor* learning_rate_out,
+                    MetaTensor* master_param_out);
+
 void SendUERecvInferMeta(const MetaTensor& x,
                          const MetaTensor& y,
                          const MetaTensor& src_index,

paddle/phi/kernels/cpu/rprop_kernel.cc

@@ -0,0 +1,143 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/rprop_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/jit/kernels.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RpropKernelCPUImpl(const Context& dev_ctx,
+                        const DenseTensor& param,
+                        const DenseTensor& grad,
+                        const DenseTensor& prev,
+                        const DenseTensor& learning_rate,
+                        const DenseTensor& learning_rate_range,
+                        const DenseTensor& etas,
+                        DenseTensor* param_out,
+                        DenseTensor* prev_out,
+                        DenseTensor* learning_rate_out) {
+  auto param_eigen = EigenVector<T>::Flatten(param);
+  auto prev_eigen = EigenVector<T>::Flatten(prev);
+  auto param_out_eigen = EigenVector<T>::Flatten(*param_out);
+  auto prev_out_eigen = EigenVector<T>::Flatten(*prev_out);
+  auto learning_rate_out_eigen = EigenVector<T>::Flatten(*learning_rate_out);
+  auto learning_rate_min = learning_rate_range.data<T>()[0];
+  auto learning_rate_max = learning_rate_range.data<T>()[1];
+  auto eta_negative = etas.data<T>()[0];
+  auto eta_positive = etas.data<T>()[1];
+
+  DenseTensor* grad_tensor = new DenseTensor();
+  grad_tensor->Resize(grad.dims());
+  dev_ctx.template Alloc<T>(grad_tensor);
+  phi::Copy<Context>(dev_ctx, grad, dev_ctx.GetPlace(), true, grad_tensor);
+  auto grad_eigen = EigenVector<T>::Flatten(*grad_tensor);
+
+  DenseTensor* product_tensor = new DenseTensor();
+  product_tensor->Resize(grad.dims());
+  dev_ctx.template Alloc<T>(product_tensor);
+  auto product_eigen = EigenVector<T>::Flatten(*product_tensor);
+
+  DenseTensor* learning_rate_tensor = new DenseTensor();
+  learning_rate_tensor->Resize(learning_rate.dims());
+  dev_ctx.template Alloc<T>(learning_rate_tensor);
+  phi::Copy<Context>(
+      dev_ctx, learning_rate, dev_ctx.GetPlace(), true, learning_rate_tensor);
+  auto learning_rate_eigen = EigenVector<T>::Flatten(*learning_rate_tensor);
+
+  DenseTensor* eta_tensor = new DenseTensor();
+  eta_tensor->Resize(learning_rate.dims());
+  dev_ctx.template Alloc<T>(eta_tensor);
+  auto eta_eigen = EigenVector<T>::Flatten(*eta_tensor);
+
+  product_eigen = grad_eigen * prev_eigen;
+  T* product_data = product_tensor->data<T>();
+  T* grad_data = grad_tensor->data<T>();
+  T* eta_data = eta_tensor->data<T>();
+  T zero = static_cast<T>(0);
+  T one = static_cast<T>(1);
+  for (int i = 0, n = product_tensor->numel(); i < n; i++) {
+    if (product_data[i] > zero) {
+      eta_data[i] = eta_positive;
+    } else if (product_data[i] == zero) {
+      eta_data[i] = one;
+    } else if (product_data[i] < zero) {
+      grad_data[i] = zero;
+      eta_data[i] = eta_negative;
+    }
+  }
+
+  learning_rate_eigen = learning_rate_eigen * eta_eigen;
+  T* learning_rate_data = learning_rate_tensor->data<T>();
+  for (int i = 0, n = learning_rate_tensor->numel(); i < n; i++) {
+    if (learning_rate_data[i] > learning_rate_max) {
+      learning_rate_data[i] = learning_rate_max;
+    } else if (learning_rate_data[i] < learning_rate_min) {
+      learning_rate_data[i] = learning_rate_min;
+    }
+  }
+
+  param_out_eigen = param_eigen - grad_eigen.sign() * learning_rate_eigen;
+  prev_out_eigen = grad_eigen;
+  learning_rate_out_eigen = learning_rate_eigen;
+  phi::Copy<Context>(dev_ctx, *grad_tensor, dev_ctx.GetPlace(), true, prev_out);
+  phi::Copy<Context>(dev_ctx,
+                     *learning_rate_tensor,
+                     dev_ctx.GetPlace(),
+                     true,
+                     learning_rate_out);
+}
+
+template <typename T, typename Context>
+void RpropKernel(const Context& dev_ctx,
+                 const DenseTensor& param,
+                 const DenseTensor& grad,
+                 const DenseTensor& prev,
+                 const DenseTensor& learning_rate,
+                 const paddle::optional<DenseTensor>& master_param UNUSED,
+                 const DenseTensor& learning_rate_range,
+                 const DenseTensor& etas,
+                 bool multi_precision UNUSED,
+                 DenseTensor* param_out,
+                 DenseTensor* prev_out,
+                 DenseTensor* learning_rate_out,
+                 DenseTensor* master_param_out UNUSED) {
+  dev_ctx.template Alloc<T>(param_out);
+  dev_ctx.template Alloc<T>(prev_out);
+  dev_ctx.template Alloc<T>(learning_rate_out);
+  RpropKernelCPUImpl<T, Context>(dev_ctx,
+                                 param,
+                                 grad,
+                                 prev,
+                                 learning_rate,
+                                 learning_rate_range,
+                                 etas,
+                                 param_out,
+                                 prev_out,
+                                 learning_rate_out);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(rprop,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::RpropKernel,
+                   phi::dtype::bfloat16,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/rprop_kernel.cu

@@ -0,0 +1,157 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/rprop_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_helper.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/mixed_vector.h"
+
+namespace phi {
+
+template <typename T, typename MT>
+__global__ void RpropKernelGPUImpl(const T* param,
+                                   const T* grad,
+                                   const T* prev,
+                                   const T* learning_rate,
+                                   const MT* master_param,
+                                   const T* learning_rate_range,
+                                   const T* etas,
+                                   int num,
+                                   T* param_out,
+                                   T* prev_out,
+                                   T* learning_rate_out,
+                                   MT* master_param_out) {
+  MT learning_rate_min_data = static_cast<MT>(learning_rate_range[0]);
+  MT learning_rate_max_data = static_cast<MT>(learning_rate_range[1]);
+  MT eta_negative_data = static_cast<MT>(etas[0]);
+  MT eta_positive_data = static_cast<MT>(etas[1]);
+  MT zero_data = static_cast<MT>(0);
+  MT one_data = static_cast<MT>(1);
+  MT negative_one_data = static_cast<MT>(-1);
+
+  CUDA_KERNEL_LOOP(i, num) {
+    MT param_data = master_param ? master_param[i] : static_cast<MT>(param[i]);
+    MT grad_data = static_cast<MT>(grad[i]);
+    MT prev_data = static_cast<MT>(prev[i]);
+    MT learning_rate_data = static_cast<MT>(learning_rate[i]);
+    MT product_data = grad_data * prev_data;
+
+    MT eta_data = one_data;
+    if (product_data > zero_data) {
+      eta_data = eta_positive_data;
+    } else if (product_data < zero_data) {
+      grad_data = zero_data;
+      eta_data = eta_negative_data;
+    }
+
+    learning_rate_data = learning_rate_data * eta_data;
+    if (learning_rate_data > learning_rate_max_data) {
+      learning_rate_data = learning_rate_max_data;
+    } else if (learning_rate_data < learning_rate_min_data) {
+      learning_rate_data = learning_rate_min_data;
+    }
+
+    MT grad_sign_data = zero_data;
+    if (grad_data > zero_data) {
+      grad_sign_data = one_data;
+    } else if (grad_data < zero_data) {
+      grad_sign_data = negative_one_data;
+    }
+
+    param_data = param_data - grad_sign_data * learning_rate_data;
+    prev_data = grad_data;
+
+    param_out[i] = static_cast<T>(param_data);
+    prev_out[i] = static_cast<T>(prev_data);
+    learning_rate_out[i] = static_cast<T>(learning_rate_data);
+    if (master_param_out) {
+      master_param_out[i] = param_data;
+    }
+  }
+}
+
+template <typename T, typename Context>
+void RpropKernel(const Context& dev_ctx,
+                 const DenseTensor& param,
+                 const DenseTensor& grad,
+                 const DenseTensor& prev,
+                 const DenseTensor& learning_rate,
+                 const paddle::optional<DenseTensor>& master_param,
+                 const DenseTensor& learning_rate_range,
+                 const DenseTensor& etas,
+                 bool multi_precision,
+                 DenseTensor* param_out,
+                 DenseTensor* prev_out,
+                 DenseTensor* learning_rate_out,
+                 DenseTensor* master_param_out) {
+  using MPDType = typename phi::dtype::MPTypeTrait<T>::Type;
+  const MPDType* master_in_data =
+      multi_precision ? master_param->data<MPDType>() : nullptr;
+  MPDType* master_out_data =
+      multi_precision ? dev_ctx.template Alloc<MPDType>(master_param_out)
+                      : nullptr;
+
+  int block = 512;
+  int grid = (param.numel() + block - 1) / block;
+
+  RpropKernelGPUImpl<T, MPDType><<<grid, block, 0, dev_ctx.stream()>>>(
+      param.data<T>(),
+      grad.data<T>(),
+      prev.data<T>(),
+      learning_rate.data<T>(),
+      master_in_data,
+      learning_rate_range.data<T>(),
+      etas.data<T>(),
+      param.numel(),
+      dev_ctx.template Alloc<T>(param_out),
+      dev_ctx.template Alloc<T>(prev_out),
+      dev_ctx.template Alloc<T>(learning_rate_out),
+      master_out_data);
+}
+
+}  // namespace phi
+
+#ifdef PADDLE_WITH_CUDA
+PD_REGISTER_KERNEL(rprop,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::RpropKernel,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   float,
+                   double) {
+  if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
+      kernel_key.dtype() == phi::DataType::BFLOAT16) {
+    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+  }
+}
+#endif
+
+#ifdef PADDLE_WITH_HIP
+PD_REGISTER_KERNEL(rprop,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::RpropKernel,
+                   phi::dtype::float16,
+                   float,
+                   double) {
+  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
+    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+  }
+}
+#endif