wehub-resource-sync
173 lines
6.8 KiB
C++
173 lines
6.8 KiB
C++
// Copyright (c) 2024 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 "paddle/phi/kernels/funcs/detail/activation_functions.h"
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/gru_compute.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/funcs/sequence2batch.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename Context, typename T>
|
|
void ReorderInitState(const Context &dev_ctx,
|
|
const DenseTensor &src,
|
|
Vector<size_t> index_lod,
|
|
DenseTensor *dst,
|
|
bool indexed_src) {
|
|
funcs::CopyMatrixRowsFunctor<Context, T> row_shuffle;
|
|
dst->Resize(src.dims());
|
|
dev_ctx.template Alloc<T>(dst);
|
|
row_shuffle(dev_ctx, src, index_lod, dst, indexed_src);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void GRUGradKernel(const Context &dev_ctx,
|
|
const DenseTensor &input,
|
|
const optional<DenseTensor> &h0_param,
|
|
const DenseTensor &weight,
|
|
const optional<DenseTensor> &bias,
|
|
const DenseTensor &batch_gate,
|
|
const DenseTensor &batch_reset_hidden_prev,
|
|
const DenseTensor &batch_hidden,
|
|
const DenseTensor &hidden,
|
|
const DenseTensor &hidden_grad,
|
|
const std::string &activation,
|
|
const std::string &gate_activation,
|
|
bool is_reverse,
|
|
bool origin_mode,
|
|
bool is_test,
|
|
DenseTensor *input_grad,
|
|
DenseTensor *h0_grad,
|
|
DenseTensor *weight_grad,
|
|
DenseTensor *bias_grad) {
|
|
auto *h0 = h0_param.get_ptr();
|
|
const T *weight_data = weight.data<T>();
|
|
|
|
auto gate_dims = batch_gate.dims();
|
|
auto hidden_dims = hidden.dims();
|
|
int frame_size = hidden_dims[1];
|
|
|
|
funcs::DenseTensor2BatchFunctor<Context, T> to_batch;
|
|
DenseTensor batch_hidden_grad, batch_gate_grad, batch_reset_hidden_prev_grad;
|
|
batch_hidden_grad.Resize(hidden_dims);
|
|
batch_gate_grad.Resize(gate_dims);
|
|
batch_reset_hidden_prev_grad.Resize(hidden_dims);
|
|
dev_ctx.template Alloc<T>(&batch_hidden_grad);
|
|
dev_ctx.template Alloc<T>(&batch_gate_grad);
|
|
dev_ctx.template Alloc<T>(&batch_reset_hidden_prev_grad);
|
|
|
|
funcs::SetConstant<Context, T> zero;
|
|
zero(dev_ctx, &batch_hidden_grad, static_cast<T>(0.0));
|
|
zero(dev_ctx, &batch_gate_grad, static_cast<T>(0.0));
|
|
zero(dev_ctx, &batch_reset_hidden_prev_grad, static_cast<T>(0.0));
|
|
|
|
DenseTensor ordered_h0, ordered_h0_grad;
|
|
|
|
Vector<size_t> order(batch_gate.lod()[2]);
|
|
|
|
if (h0) {
|
|
ReorderInitState<Context, T>(dev_ctx, *h0, order, &ordered_h0, true);
|
|
}
|
|
if (h0_grad) {
|
|
ordered_h0_grad.Resize(h0_grad->dims());
|
|
dev_ctx.template Alloc<T>(&ordered_h0_grad);
|
|
zero(dev_ctx, &ordered_h0_grad, static_cast<T>(0.0));
|
|
}
|
|
|
|
batch_hidden_grad.set_lod(batch_hidden.lod());
|
|
to_batch(dev_ctx, hidden_grad, &batch_hidden_grad, false, is_reverse);
|
|
|
|
funcs::GRUMetaValue<T> gru_value;
|
|
gru_value.gate_weight = const_cast<T *>(weight_data);
|
|
gru_value.state_weight =
|
|
const_cast<T *>(weight_data + 2 * frame_size * frame_size);
|
|
|
|
funcs::GRUMetaGrad<T> gru_grad;
|
|
if (weight_grad) {
|
|
gru_grad.gate_weight_grad = dev_ctx.template Alloc<T>(weight_grad);
|
|
zero(dev_ctx, weight_grad, static_cast<T>(0.0));
|
|
gru_grad.state_weight_grad =
|
|
weight_grad->data<T>() + 2 * frame_size * frame_size;
|
|
} else {
|
|
gru_grad.gate_weight_grad = nullptr;
|
|
gru_grad.state_weight_grad = nullptr;
|
|
}
|
|
|
|
auto batch_starts = batch_hidden_grad.lod()[0];
|
|
size_t num_batch = batch_starts.size() - 1;
|
|
auto active_node = funcs::detail::GetActivationType(activation);
|
|
auto active_gate = funcs::detail::GetActivationType(gate_activation);
|
|
for (int n = static_cast<int>(num_batch) - 1; n >= 0; n--) {
|
|
int bstart = static_cast<int>(batch_starts[n]);
|
|
int bend = static_cast<int>(batch_starts[n + 1]);
|
|
int cur_batch_size = bend - bstart;
|
|
|
|
DenseTensor gate_t = batch_gate.Slice(bstart, bend);
|
|
gru_value.gate_value = gate_t.data<T>();
|
|
DenseTensor reset_hidden_prev_t =
|
|
batch_reset_hidden_prev.Slice(bstart, bend);
|
|
gru_value.reset_output_value = reset_hidden_prev_t.data<T>();
|
|
|
|
DenseTensor hidden_grad_t = batch_hidden_grad.Slice(bstart, bend);
|
|
gru_grad.output_grad = hidden_grad_t.data<T>();
|
|
DenseTensor gate_grad_t = batch_gate_grad.Slice(bstart, bend);
|
|
gru_grad.gate_grad = gate_grad_t.data<T>();
|
|
DenseTensor reset_hidden_prev_grad_t =
|
|
batch_reset_hidden_prev_grad.Slice(bstart, bend);
|
|
gru_grad.reset_output_grad = reset_hidden_prev_grad_t.data<T>();
|
|
if (n == 0) {
|
|
gru_value.prev_out_value = h0 ? ordered_h0.data<T>() : nullptr;
|
|
gru_grad.prev_out_grad =
|
|
h0 && h0_grad ? ordered_h0_grad.data<T>() : nullptr;
|
|
} else {
|
|
int bstart_pre = static_cast<int>(batch_starts[n - 1]);
|
|
DenseTensor hidden_prev_t = batch_hidden.Slice(bstart_pre, bstart);
|
|
gru_value.prev_out_value = hidden_prev_t.data<T>();
|
|
DenseTensor hidden_prev_grad_t =
|
|
batch_hidden_grad.Slice(bstart_pre, bstart);
|
|
gru_grad.prev_out_grad = hidden_prev_grad_t.data<T>();
|
|
}
|
|
gru_value.output_value = nullptr;
|
|
funcs::GRUUnitGradFunctor<Context, T>::compute(dev_ctx,
|
|
gru_value,
|
|
gru_grad,
|
|
frame_size,
|
|
cur_batch_size,
|
|
active_node,
|
|
active_gate,
|
|
origin_mode);
|
|
}
|
|
if (input_grad) {
|
|
dev_ctx.template Alloc<T>(input_grad);
|
|
funcs::Batch2DenseTensorFunctor<Context, T> to_seq;
|
|
batch_gate_grad.set_lod(batch_gate.lod());
|
|
to_seq(dev_ctx, batch_gate_grad, input_grad);
|
|
}
|
|
if (bias_grad) {
|
|
dev_ctx.template Alloc<T>(bias_grad);
|
|
funcs::ColwiseSum<Context, T> col_sum;
|
|
col_sum(dev_ctx, batch_gate_grad, bias_grad);
|
|
}
|
|
if (h0_param && h0_grad) {
|
|
ReorderInitState<Context, T>(
|
|
dev_ctx, ordered_h0_grad, order, h0_grad, false);
|
|
}
|
|
}
|
|
} // namespace phi
|