wehub-resource-sync
324 lines
11 KiB
C++
324 lines
11 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include <math.h>
|
|
|
|
#include "paddle/phi/common/amp_type_traits.h"
|
|
#include "paddle/phi/kernels/funcs/algorithm.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/for_range.h"
|
|
#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
|
|
#include "paddle/phi/kernels/rmsprop_kernel.h"
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
struct RmsFunctor {
|
|
RmsFunctor(const Context &dev_ctx,
|
|
const DenseTensor ¶m,
|
|
const DenseTensor &mean_square,
|
|
const DenseTensor &grad,
|
|
const DenseTensor &moment,
|
|
const DenseTensor &learning_rate,
|
|
const optional<DenseTensor> &mean_grad_opt,
|
|
const optional<DenseTensor> &master_param,
|
|
float epsilon_t,
|
|
float decay_t,
|
|
float momentum_t,
|
|
bool centered,
|
|
bool multi_precision,
|
|
DenseTensor *param_out,
|
|
DenseTensor *moment_out,
|
|
DenseTensor *mean_square_out,
|
|
DenseTensor *mean_grad_out,
|
|
DenseTensor *master_param_outs);
|
|
};
|
|
|
|
template <typename T>
|
|
struct DenseRmspropGradFunctor {
|
|
inline explicit DenseRmspropGradFunctor(const T *grad) : grad_(grad) {}
|
|
|
|
HOSTDEVICE inline T operator()(int64_t idx) const { return grad_[idx]; }
|
|
|
|
const T *grad_;
|
|
};
|
|
|
|
template <typename T>
|
|
struct SparseRmspropGradFunctor {
|
|
inline SparseRmspropGradFunctor(const T *grad,
|
|
const int64_t *rows,
|
|
int64_t row_numel,
|
|
int64_t row_count)
|
|
: grad_(grad),
|
|
rows_(rows),
|
|
row_numel_(row_numel),
|
|
row_count_(row_count) {}
|
|
|
|
HOSTDEVICE inline T operator()(int64_t idx) const {
|
|
auto row_idx = funcs::BinarySearch(rows_, row_count_, idx / row_numel_);
|
|
return row_idx >= 0 ? grad_[row_idx * row_numel_ + idx % row_numel_]
|
|
: static_cast<T>(0);
|
|
}
|
|
|
|
const T *grad_;
|
|
const int64_t *rows_;
|
|
int64_t row_numel_;
|
|
int64_t row_count_;
|
|
};
|
|
|
|
template <typename T, typename MT, typename GradFunctor>
|
|
struct UncenteredRmspropFunctor {
|
|
UncenteredRmspropFunctor(T *param,
|
|
MT *ms,
|
|
MT *mom,
|
|
const MT *lr,
|
|
MT *master_p,
|
|
MT rho,
|
|
MT epsilon,
|
|
MT momentum,
|
|
const GradFunctor &grad_functor)
|
|
: param_(param),
|
|
ms_(ms),
|
|
mom_(mom),
|
|
master_p_(master_p),
|
|
lr_(lr),
|
|
rho_(rho),
|
|
epsilon_(epsilon),
|
|
momentum_(momentum),
|
|
grad_functor_(grad_functor) {}
|
|
|
|
HOSTDEVICE inline void operator()(int64_t idx) const {
|
|
MT g = static_cast<MT>(grad_functor_(idx));
|
|
MT l_rho = static_cast<MT>(1) - rho_;
|
|
MT ms_out = rho_ * ms_[idx] + l_rho * g * g;
|
|
MT mom_out = momentum_ * mom_[idx] +
|
|
static_cast<MT>(lr_[0]) * g / sqrt(ms_out + epsilon_);
|
|
MT p = master_p_ ? master_p_[idx] : static_cast<MT>(param_[idx]);
|
|
MT p_m = p - mom_out;
|
|
param_[idx] = static_cast<T>(p_m);
|
|
ms_[idx] = ms_out;
|
|
mom_[idx] = mom_out;
|
|
if (master_p_) master_p_[idx] = p_m;
|
|
}
|
|
|
|
T *param_;
|
|
MT *ms_;
|
|
MT *mom_;
|
|
MT *master_p_;
|
|
const MT *lr_;
|
|
MT rho_;
|
|
MT epsilon_;
|
|
MT momentum_;
|
|
GradFunctor grad_functor_;
|
|
};
|
|
|
|
template <typename T, typename MT, typename GradFunctor>
|
|
struct CenteredRmspropFunctor {
|
|
CenteredRmspropFunctor(T *param,
|
|
MT *ms,
|
|
MT *mom,
|
|
MT *mean_grad,
|
|
const MT *lr,
|
|
MT *master_param,
|
|
MT rho,
|
|
MT epsilon,
|
|
MT momentum,
|
|
const GradFunctor &grad_functor)
|
|
: param_(param),
|
|
ms_(ms),
|
|
mom_(mom),
|
|
master_p_(master_param),
|
|
mean_grad_(mean_grad),
|
|
lr_(lr),
|
|
rho_(rho),
|
|
epsilon_(epsilon),
|
|
momentum_(momentum),
|
|
grad_functor_(grad_functor) {}
|
|
|
|
HOSTDEVICE inline void operator()(int64_t idx) const {
|
|
MT g = static_cast<MT>(grad_functor_(idx));
|
|
MT l_rho = static_cast<MT>(1) - rho_;
|
|
MT ms_out = rho_ * ms_[idx] + l_rho * g * g;
|
|
MT mg_out = rho_ * mean_grad_[idx] + l_rho * g;
|
|
MT mom_out =
|
|
momentum_ * mom_[idx] +
|
|
static_cast<MT>(lr_[0]) * g / sqrt(ms_out - mg_out * mg_out + epsilon_);
|
|
|
|
MT p = master_p_ ? master_p_[idx] : static_cast<MT>(param_[idx]);
|
|
MT p_m = p - mom_out;
|
|
param_[idx] = static_cast<T>(p_m);
|
|
ms_[idx] = ms_out;
|
|
mom_[idx] = mom_out;
|
|
mean_grad_[idx] = mg_out;
|
|
if (master_p_) master_p_[idx] = p_m;
|
|
}
|
|
|
|
T *param_;
|
|
MT *ms_;
|
|
MT *mom_;
|
|
MT *master_p_;
|
|
MT *mean_grad_;
|
|
const MT *lr_;
|
|
MT rho_;
|
|
MT epsilon_;
|
|
MT momentum_;
|
|
GradFunctor grad_functor_;
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void RmspropDenseKernel(const Context &dev_ctx,
|
|
const DenseTensor ¶m,
|
|
const DenseTensor &mean_square,
|
|
const DenseTensor &grad,
|
|
const DenseTensor &moment,
|
|
const DenseTensor &learning_rate,
|
|
const optional<DenseTensor> &mean_grad_opt,
|
|
const optional<DenseTensor> &master_param,
|
|
float epsilon_t,
|
|
float decay_t,
|
|
float momentum_t,
|
|
bool centered,
|
|
bool multi_precision,
|
|
DenseTensor *param_out,
|
|
DenseTensor *moment_out,
|
|
DenseTensor *mean_square_out,
|
|
DenseTensor *mean_grad_out,
|
|
DenseTensor *master_param_outs) {
|
|
RmsFunctor<T, Context> functor(dev_ctx,
|
|
param,
|
|
mean_square,
|
|
grad,
|
|
moment,
|
|
learning_rate,
|
|
mean_grad_opt,
|
|
master_param,
|
|
epsilon_t,
|
|
decay_t,
|
|
momentum_t,
|
|
centered,
|
|
multi_precision,
|
|
param_out,
|
|
moment_out,
|
|
mean_square_out,
|
|
mean_grad_out,
|
|
master_param_outs);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void RmspropSparseKernel(const Context &dev_ctx,
|
|
const DenseTensor ¶m,
|
|
const DenseTensor &mean_square,
|
|
const SelectedRows &grad,
|
|
const DenseTensor &moment,
|
|
const DenseTensor &learning_rate,
|
|
const optional<DenseTensor> &mean_grad_opt,
|
|
const optional<DenseTensor> &master_param UNUSED,
|
|
float epsilon_t,
|
|
float decay_t,
|
|
float momentum_t,
|
|
bool centered,
|
|
bool multi_precision,
|
|
DenseTensor *param_out,
|
|
DenseTensor *moment_out,
|
|
DenseTensor *mean_square_out,
|
|
DenseTensor *mean_grad_out,
|
|
DenseTensor *master_param_outs) {
|
|
using MT = typename MPTypeTrait<T>::Type;
|
|
auto epsilon = static_cast<MT>(epsilon_t);
|
|
auto rho = static_cast<MT>(decay_t);
|
|
auto momentum = static_cast<MT>(momentum_t);
|
|
|
|
auto &p_tensor = param;
|
|
auto &ms_tensor = mean_square;
|
|
auto &lr_tensor = learning_rate;
|
|
auto &mom_tensor = moment;
|
|
|
|
PADDLE_ENFORCE_EQ(p_tensor.IsSharedBufferWith(*param_out),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Param and ParamOut must be the same Tensor"));
|
|
PADDLE_ENFORCE_EQ(mom_tensor.IsSharedBufferWith(*moment_out),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Moment and MomentOut must be the same Tensor"));
|
|
PADDLE_ENFORCE_EQ(
|
|
ms_tensor.IsSharedBufferWith(*mean_square_out),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"MeanSquare and MeanSquareOut must be the same Tensor"));
|
|
size_t limit = static_cast<size_t>(ms_tensor.numel());
|
|
|
|
SelectedRows tmp_merged_grad;
|
|
SelectedRows *merged_grad = &tmp_merged_grad;
|
|
funcs::scatter::MergeAdd<Context, T> merge_func;
|
|
merge_func(dev_ctx, grad, merged_grad);
|
|
|
|
funcs::ForRange<Context> for_range(dev_ctx, limit);
|
|
auto &grad_merge_rows = merged_grad->rows();
|
|
MixVector<int64_t> mixv_grad_merge_rows(&grad_merge_rows);
|
|
const int64_t *rows = mixv_grad_merge_rows.Data(dev_ctx.GetPlace());
|
|
|
|
auto &merged_tensor = merged_grad->value();
|
|
int64_t row_count = merged_grad->rows().size();
|
|
int64_t row_numel = merged_tensor.numel() / row_count;
|
|
SparseRmspropGradFunctor<T> grad_func(
|
|
merged_tensor.data<T>(), rows, row_numel, row_count);
|
|
|
|
MT *master_out_data =
|
|
multi_precision ? dev_ctx.template Alloc<MT>(master_param_outs) : nullptr;
|
|
|
|
if (centered) {
|
|
auto mg_tensor = mean_grad_opt.get_ptr();
|
|
if (mg_tensor) {
|
|
PADDLE_ENFORCE_EQ(
|
|
mg_tensor->Holder(),
|
|
mean_grad_out->Holder(),
|
|
common::errors::InvalidArgument(
|
|
"MeanGrad and MeanGradOut must be the same Tensor"));
|
|
} else {
|
|
PADDLE_ENFORCE_EQ(
|
|
mg_tensor,
|
|
mean_grad_out,
|
|
common::errors::InvalidArgument(
|
|
"MeanGrad and MeanGradOut must be the same Tensor"));
|
|
}
|
|
|
|
for_range(CenteredRmspropFunctor<T, MT, SparseRmspropGradFunctor<T>>(
|
|
dev_ctx.template Alloc<T>(param_out),
|
|
dev_ctx.template Alloc<MT>(mean_square_out),
|
|
dev_ctx.template Alloc<MT>(moment_out),
|
|
dev_ctx.template Alloc<MT>(mean_grad_out),
|
|
lr_tensor.data<MT>(),
|
|
master_out_data,
|
|
rho,
|
|
epsilon,
|
|
momentum,
|
|
grad_func));
|
|
} else {
|
|
for_range(UncenteredRmspropFunctor<T, MT, SparseRmspropGradFunctor<T>>(
|
|
dev_ctx.template Alloc<T>(param_out),
|
|
dev_ctx.template Alloc<MT>(mean_square_out),
|
|
dev_ctx.template Alloc<MT>(moment_out),
|
|
lr_tensor.data<MT>(),
|
|
master_out_data,
|
|
rho,
|
|
epsilon,
|
|
momentum,
|
|
grad_func));
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|