wehub-resource-sync
339 lines
11 KiB
C++
339 lines
11 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 <memory>
|
|
#include "paddle/phi/backends/all_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/activation_functor.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/utils/optional.h"
|
|
namespace phi {
|
|
|
|
enum GRUActivationType { identity = 0, sigmoid = 1, tanh = 2, relu = 3 };
|
|
|
|
template <typename T, typename Device, typename X, typename Y>
|
|
void ActCompute(const int act_type, const Device& d, X x, Y y, Place place) {
|
|
if (act_type == identity) {
|
|
y.device(d) = x;
|
|
} else if (act_type == sigmoid) {
|
|
funcs::SigmoidFunctor<T>()(d, x, y);
|
|
} else if (act_type == tanh) {
|
|
funcs::TanhFunctor<T>()(d, x, y);
|
|
} else if (act_type == relu) {
|
|
if (place == CPUPlace())
|
|
funcs::ReluCPUFunctor<T>()(d, x, y);
|
|
else
|
|
funcs::ReluCUDAFunctor<T>()(d, x, y);
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Unsupported activation type, only supports identity, sigmoid, tanh "
|
|
"and relu."));
|
|
}
|
|
}
|
|
|
|
#define ACT_COMPUTE ActCompute<T>
|
|
|
|
template <typename T, typename Context>
|
|
void GRUUnitKernel(const Context& dev_ctx,
|
|
const DenseTensor& input,
|
|
const DenseTensor& hidden_prev,
|
|
const DenseTensor& weight,
|
|
const optional<DenseTensor>& bias,
|
|
int activation,
|
|
int gate_activation,
|
|
bool origin_mode,
|
|
DenseTensor* gate,
|
|
DenseTensor* reset_hidden_prev,
|
|
DenseTensor* hidden) {
|
|
auto* input_p = &input;
|
|
auto* hidden_prev_p = &hidden_prev;
|
|
|
|
dev_ctx.template Alloc<T>(gate);
|
|
dev_ctx.template Alloc<T>(reset_hidden_prev);
|
|
dev_ctx.template Alloc<T>(hidden);
|
|
|
|
int64_t batch_size = input_p->dims()[0];
|
|
int64_t frame_size = hidden_prev_p->dims()[1];
|
|
|
|
auto x = EigenMatrix<T>::From(input);
|
|
auto h_p = EigenMatrix<T>::From(hidden_prev);
|
|
auto g = EigenMatrix<T>::From(*gate);
|
|
auto r_h_p = EigenMatrix<T>::From(*reset_hidden_prev);
|
|
auto h = EigenMatrix<T>::From(*hidden);
|
|
auto& place = *dev_ctx.eigen_device();
|
|
|
|
// calculate unactivated gate outputs
|
|
if (bias) {
|
|
auto b = EigenMatrix<T>::From(bias.get());
|
|
g.device(place) =
|
|
x + b.reshape(Eigen::array<int64_t, 2>({{1, frame_size * 3}}))
|
|
.broadcast(Eigen::array<int64_t, 2>({{batch_size, 1}}));
|
|
} else {
|
|
g.device(place) = x;
|
|
}
|
|
const T* hidden_prev_data = hidden_prev.data<T>();
|
|
const T* weight_data = weight.data<T>();
|
|
T* gate_data = gate->data<T>();
|
|
T* reset_hidden_prev_data = reset_hidden_prev->data<T>();
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
blas.GEMM(false,
|
|
false,
|
|
batch_size,
|
|
2 * frame_size,
|
|
frame_size,
|
|
1,
|
|
hidden_prev_data,
|
|
frame_size,
|
|
weight_data,
|
|
frame_size * 2,
|
|
1,
|
|
gate_data,
|
|
frame_size * 3);
|
|
|
|
// calculate activated gate
|
|
Eigen::array<int64_t, 2> extents{{batch_size, frame_size}};
|
|
Eigen::array<int64_t, 2> u_offsets{{0, 0}};
|
|
ACT_COMPUTE(gate_activation,
|
|
place,
|
|
g.slice(u_offsets, extents),
|
|
g.slice(u_offsets, extents),
|
|
dev_ctx.GetPlace());
|
|
auto u = g.slice(u_offsets, extents); // update gate
|
|
Eigen::array<int64_t, 2> r_offsets{{0, frame_size}};
|
|
ACT_COMPUTE(gate_activation,
|
|
place,
|
|
g.slice(r_offsets, extents),
|
|
g.slice(r_offsets, extents),
|
|
dev_ctx.GetPlace());
|
|
auto r = g.slice(r_offsets, extents); // reset gate
|
|
r_h_p.device(place) = r * h_p; // reset previous hidden state
|
|
blas.GEMM(false,
|
|
false,
|
|
batch_size,
|
|
frame_size,
|
|
frame_size,
|
|
1,
|
|
reset_hidden_prev_data,
|
|
frame_size,
|
|
weight_data + frame_size * frame_size * 2,
|
|
frame_size,
|
|
1,
|
|
gate_data + frame_size * 2,
|
|
frame_size * 3);
|
|
|
|
Eigen::array<int64_t, 2> c_offsets{{0, frame_size * 2}};
|
|
ACT_COMPUTE(activation,
|
|
place,
|
|
g.slice(c_offsets, extents),
|
|
g.slice(c_offsets, extents),
|
|
dev_ctx.GetPlace());
|
|
auto c = g.slice(c_offsets, extents); // output candidate
|
|
|
|
// calculate final output
|
|
if (origin_mode) {
|
|
h.device(place) = c + u * (h_p - c); // (1 - u) * c + u * h_p
|
|
} else {
|
|
h.device(place) = u * (c - h_p) + h_p; // u * c + (1 - u) * h_p
|
|
}
|
|
}
|
|
|
|
template <typename T,
|
|
typename Device,
|
|
typename X,
|
|
typename Y,
|
|
typename DX,
|
|
typename DY>
|
|
void ActGradCompute(
|
|
const int act_type, const Device& d, X x, Y y, DX dx, DY dy) {
|
|
// x is dummy and won't be used even in Relu(use y instead)
|
|
if (act_type == identity)
|
|
dx.device(d) = dy;
|
|
else if (act_type == sigmoid)
|
|
funcs::SigmoidGradFunctor<T>()(d, x, y, dy, dx);
|
|
else if (act_type == tanh)
|
|
funcs::TanhGradFunctor<T>()(d, x, y, dy, dx);
|
|
else if (act_type == relu)
|
|
funcs::ReluGradFunctor<T>()(d, x, y, dy, dx);
|
|
else
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Unsupported activation type, only supports identity, sigmoid, tanh "
|
|
"and relu."));
|
|
}
|
|
|
|
#define ACT_GRAD_COMPUTE ActGradCompute<T>
|
|
|
|
template <typename T, typename Context>
|
|
void GRUUnitGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& input,
|
|
const DenseTensor& hidden_prev,
|
|
const DenseTensor& weight,
|
|
const optional<DenseTensor>& bias,
|
|
const DenseTensor& gate,
|
|
const DenseTensor& reset_hidden_prev,
|
|
const DenseTensor& hidden_grad,
|
|
int activation,
|
|
int gate_activation,
|
|
bool origin_mode,
|
|
DenseTensor* input_grad,
|
|
DenseTensor* hidden_prev_grad,
|
|
DenseTensor* weight_grad,
|
|
DenseTensor* bias_grad) {
|
|
DenseTensor gate_grad;
|
|
DenseTensor reset_hidden_prev_grad;
|
|
|
|
const T* hidden_prev_data = hidden_prev.data<T>();
|
|
const T* weight_data = weight.data<T>();
|
|
gate_grad.Resize(input.dims());
|
|
T* gate_grad_data = dev_ctx.template Alloc<T>(&gate_grad);
|
|
const T* reset_hidden_prev_data = reset_hidden_prev.data<T>();
|
|
reset_hidden_prev_grad.Resize(reset_hidden_prev.dims());
|
|
T* reset_hidden_prev_grad_data =
|
|
dev_ctx.template Alloc<T>(&reset_hidden_prev_grad);
|
|
|
|
auto h_p = EigenMatrix<T>::From(hidden_prev);
|
|
auto g = EigenMatrix<T>::From(gate);
|
|
auto d_h = EigenMatrix<T>::From(hidden_grad);
|
|
auto d_g = EigenMatrix<T>::From(gate_grad);
|
|
auto d_r_h_p = EigenMatrix<T>::From(reset_hidden_prev_grad);
|
|
auto& place = *dev_ctx.eigen_device();
|
|
|
|
int64_t batch_size = input.dims()[0];
|
|
int64_t frame_size = hidden_prev.dims()[1];
|
|
|
|
Eigen::array<int64_t, 2> extents{{batch_size, frame_size}};
|
|
Eigen::array<int64_t, 2> u_offsets{{0, 0}};
|
|
auto u = g.slice(u_offsets, extents); // update gate
|
|
Eigen::array<int64_t, 2> r_offsets{{0, frame_size}};
|
|
auto r = g.slice(r_offsets, extents); // reset gate
|
|
Eigen::array<int64_t, 2> c_offsets{{0, frame_size * 2}};
|
|
auto c = g.slice(c_offsets, extents); // output candidate
|
|
|
|
// backward for unactivated update gate
|
|
if (origin_mode) {
|
|
ACT_GRAD_COMPUTE(gate_activation,
|
|
place,
|
|
u,
|
|
u,
|
|
d_g.slice(u_offsets, extents),
|
|
d_h * (h_p - c));
|
|
// backward for unactivated output candidate
|
|
ACT_GRAD_COMPUTE(
|
|
activation, place, c, c, d_g.slice(c_offsets, extents), d_h * (1 - u));
|
|
} else {
|
|
ACT_GRAD_COMPUTE(gate_activation,
|
|
place,
|
|
u,
|
|
u,
|
|
d_g.slice(u_offsets, extents),
|
|
d_h * (c - h_p));
|
|
// backward for unactivated output candidate
|
|
ACT_GRAD_COMPUTE(
|
|
activation, place, c, c, d_g.slice(c_offsets, extents), d_h * u);
|
|
}
|
|
// backward for reset_hidden_prev
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
blas.GEMM(false,
|
|
true,
|
|
batch_size,
|
|
frame_size,
|
|
frame_size,
|
|
1,
|
|
gate_grad_data + frame_size * 2,
|
|
frame_size * 3,
|
|
weight_data + frame_size * frame_size * 2,
|
|
frame_size,
|
|
0,
|
|
reset_hidden_prev_grad_data,
|
|
frame_size);
|
|
// backward for unactivated reset gate
|
|
ACT_GRAD_COMPUTE(gate_activation,
|
|
place,
|
|
r,
|
|
r,
|
|
d_g.slice(r_offsets, extents),
|
|
d_r_h_p * h_p);
|
|
// backward for weight
|
|
if (weight_grad) {
|
|
T* weight_grad_data = dev_ctx.template Alloc<T>(weight_grad);
|
|
// backward for state_weight
|
|
blas.GEMM(true,
|
|
false,
|
|
frame_size,
|
|
frame_size,
|
|
batch_size,
|
|
1,
|
|
reset_hidden_prev_data,
|
|
frame_size,
|
|
gate_grad_data + frame_size * 2,
|
|
frame_size * 3,
|
|
0,
|
|
weight_grad_data + frame_size * frame_size * 2,
|
|
frame_size);
|
|
|
|
// backward for update_gate_weight and reset_gate_weight
|
|
blas.GEMM(true,
|
|
false,
|
|
frame_size,
|
|
frame_size * 2,
|
|
batch_size,
|
|
1,
|
|
hidden_prev_data,
|
|
frame_size,
|
|
gate_grad_data,
|
|
frame_size * 3,
|
|
0,
|
|
weight_grad_data,
|
|
frame_size * 2);
|
|
}
|
|
// backward for hidden_prev
|
|
if (hidden_prev_grad) {
|
|
T* hidden_prev_grad_data = dev_ctx.template Alloc<T>(hidden_prev_grad);
|
|
auto d_h_p = EigenMatrix<T>::From(*hidden_prev_grad);
|
|
if (origin_mode) {
|
|
d_h_p.device(place) = d_r_h_p * r + d_h * u;
|
|
} else {
|
|
d_h_p.device(place) = d_r_h_p * r + d_h * (1 - u);
|
|
}
|
|
blas.GEMM(false,
|
|
true,
|
|
batch_size,
|
|
frame_size,
|
|
frame_size * 2,
|
|
1,
|
|
gate_grad_data,
|
|
frame_size * 3,
|
|
weight_data,
|
|
frame_size * 2,
|
|
1,
|
|
hidden_prev_grad_data,
|
|
frame_size);
|
|
}
|
|
// backward for input
|
|
if (input_grad) {
|
|
dev_ctx.template Alloc<T>(input_grad);
|
|
auto d_x = EigenMatrix<T>::From(*input_grad);
|
|
d_x.device(place) = d_g;
|
|
}
|
|
// backward for bias
|
|
if (bias_grad) {
|
|
dev_ctx.template Alloc<T>(bias_grad);
|
|
auto d_b = EigenVector<T>::Flatten(*bias_grad);
|
|
d_b.device(place) = d_g.sum(Eigen::array<int64_t, 1>({{0}}));
|
|
}
|
|
}
|
|
} // namespace phi
|