#ifndef CAFFE2_OPERATORS_LSTM_UNIT_OP_H_
|
#define CAFFE2_OPERATORS_LSTM_UNIT_OP_H_
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/operator.h"
|
#include "caffe2/utils/conversions.h"
|
|
namespace caffe2 {
|
namespace detail {
|
template <typename T>
|
inline T sigmoid(T x) {
|
return 1. / (1. + exp(-x));
|
}
|
|
template <typename T>
|
inline T host_tanh(T x) {
|
return 2. * sigmoid(2. * x) - 1.;
|
}
|
|
template <typename T, typename Context>
|
void LSTMUnit(
|
int N,
|
int D,
|
int t,
|
const T* H_prev,
|
const T* C_prev,
|
const T* X,
|
const int32_t* seqLengths,
|
bool drop_states,
|
T* C,
|
T* H,
|
const float forget_bias,
|
Context* /*context*/) {
|
for (int n = 0; n < N; ++n) {
|
const bool valid = seqLengths == nullptr || t < seqLengths[n];
|
|
for (int d = 0; d < D; ++d) {
|
if (!valid) {
|
if (drop_states) {
|
H[d] = 0;
|
C[d] = 0;
|
} else {
|
H[d] = H_prev[d];
|
C[d] = C_prev[d];
|
}
|
} else {
|
const T i = sigmoid(X[d]);
|
const T f = sigmoid(X[1 * D + d] + convert::To<float, T>(forget_bias));
|
const T o = sigmoid(X[2 * D + d]);
|
const T g = host_tanh(X[3 * D + d]);
|
const T c_prev = C_prev[d];
|
const T c = f * c_prev + i * g;
|
C[d] = c;
|
const T host_tanh_c = host_tanh(c);
|
H[d] = o * host_tanh_c;
|
}
|
}
|
H_prev += D;
|
C_prev += D;
|
X += 4 * D;
|
C += D;
|
H += D;
|
}
|
}
|
|
template <typename T, typename Context>
|
void LSTMUnitGradient(
|
int N,
|
int D,
|
int t,
|
const T* C_prev,
|
const T* X,
|
const int32_t* seqLengths,
|
const T* C,
|
const T* H,
|
const T* C_diff,
|
const T* H_diff,
|
bool drop_states,
|
T* H_prev_diff,
|
T* C_prev_diff,
|
T* X_diff,
|
const float forget_bias,
|
Context* /*context*/) {
|
for (int n = 0; n < N; ++n) {
|
const bool valid = seqLengths == nullptr || t < seqLengths[n];
|
|
for (int d = 0; d < D; ++d) {
|
T* c_prev_diff = C_prev_diff + d;
|
T* h_prev_diff = H_prev_diff + d;
|
T* i_diff = X_diff + d;
|
T* f_diff = X_diff + 1 * D + d;
|
T* o_diff = X_diff + 2 * D + d;
|
T* g_diff = X_diff + 3 * D + d;
|
|
if (!valid) {
|
if (drop_states) {
|
*h_prev_diff = 0;
|
*c_prev_diff = 0;
|
} else {
|
*h_prev_diff = H_diff[d];
|
*c_prev_diff = C_diff[d];
|
}
|
*i_diff = 0;
|
*f_diff = 0;
|
*o_diff = 0;
|
*g_diff = 0;
|
} else {
|
const T i = sigmoid(X[d]);
|
const T f = sigmoid(X[1 * D + d] + convert::To<float, T>(forget_bias));
|
const T o = sigmoid(X[2 * D + d]);
|
const T g = host_tanh(X[3 * D + d]);
|
const T c_prev = C_prev[d];
|
const T c = C[d];
|
const T host_tanh_c = host_tanh(c);
|
const T c_term_diff = C_diff[d] + H_diff[d] * o * (1 - host_tanh_c * host_tanh_c);
|
*c_prev_diff = c_term_diff * f;
|
*h_prev_diff = 0; // not used in 'valid' case
|
*i_diff = c_term_diff * g * i * (1 - i);
|
*f_diff = c_term_diff * c_prev * f * (1 - f);
|
*o_diff = H_diff[d] * host_tanh_c * o * (1 - o);
|
*g_diff = c_term_diff * i * (1 - g * g);
|
}
|
}
|
C_prev += D;
|
X += 4 * D;
|
C += D;
|
H += D;
|
C_diff += D;
|
H_diff += D;
|
X_diff += 4 * D;
|
H_prev_diff += D;
|
C_prev_diff += D;
|
}
|
}
|
} // namespace detail
|
|
template <typename Context>
|
class LSTMUnitOp : public Operator<Context> {
|
public:
|
explicit LSTMUnitOp(const OperatorDef& operator_def, Workspace* ws)
|
: Operator<Context>(operator_def, ws),
|
forget_bias_(static_cast<float>(
|
this->template GetSingleArgument<float>("forget_bias", 0.0))),
|
sequence_lengths_(
|
this->template GetSingleArgument<bool>("sequence_lengths", true)),
|
drop_states_(
|
this->template GetSingleArgument<bool>("drop_states", false)) {}
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
using Operator<Context>::Operator;
|
|
template <typename T>
|
bool DoRunWithType() {
|
// handle potentially-missing sequence lengths input
|
const size_t TIMESTEP = SEQ_LENGTHS + (sequence_lengths_ ? 1 : 0);
|
|
// Extract N
|
const auto N = Input(CELL_T_M_1).size(1);
|
|
// Gates: 1xNxG
|
const auto G = Input(GATES).size(2);
|
const auto D = Input(CELL_T_M_1).size(2);
|
|
CAFFE_ENFORCE_EQ(4 * D, G);
|
const auto* H_prev = Input(HIDDEN_T_M_1).template data<T>();
|
const auto* C_prev = Input(CELL_T_M_1).template data<T>();
|
const auto* X = Input(GATES).template data<T>();
|
|
const int32_t* seqLengths = nullptr;
|
if (sequence_lengths_) {
|
CAFFE_ENFORCE_EQ(Input(SEQ_LENGTHS).numel(), N);
|
seqLengths = Input(SEQ_LENGTHS).template data<int32_t>();
|
}
|
|
const auto t = static_cast<OperatorBase*>(this)
|
->Input<Tensor>(TIMESTEP, CPU)
|
.template data<int32_t>()[0];
|
Output(CELL_T)->ResizeLike(Input(CELL_T_M_1));
|
auto* C = Output(CELL_T)->template mutable_data<T>();
|
Output(HIDDEN_T)->ResizeLike(Input(CELL_T_M_1));
|
auto* H = Output(HIDDEN_T)->template mutable_data<T>();
|
detail::LSTMUnit<T, Context>(
|
N,
|
D,
|
t,
|
H_prev,
|
C_prev,
|
X,
|
seqLengths,
|
drop_states_,
|
C,
|
H,
|
forget_bias_,
|
&context_);
|
return true;
|
}
|
|
bool RunOnDevice() override {
|
return DoRunWithType<float>();
|
}
|
|
protected:
|
INPUT_TAGS(HIDDEN_T_M_1, CELL_T_M_1, GATES, SEQ_LENGTHS);
|
// additional input tags are determined dynamically based on whether
|
// sequence_lengths is present.
|
OUTPUT_TAGS(HIDDEN_T, CELL_T);
|
|
float forget_bias_;
|
bool sequence_lengths_;
|
|
private:
|
bool drop_states_;
|
};
|
|
template <typename Context>
|
class LSTMUnitGradientOp : public Operator<Context> {
|
public:
|
template <class... Args>
|
explicit LSTMUnitGradientOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
forget_bias_(static_cast<float>(
|
this->template GetSingleArgument<float>("forget_bias", 0.0))),
|
sequence_lengths_(
|
this->template GetSingleArgument<bool>("sequence_lengths", true)),
|
drop_states_(
|
this->template GetSingleArgument<bool>("drop_states", false)) {}
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
template <typename T>
|
bool DoRunWithType() {
|
// handle potentially-missing sequence lengths input
|
const size_t inputOffset = SEQ_LENGTHS + (sequence_lengths_ ? 1 : 0);
|
const size_t TIMESTEP = inputOffset;
|
const size_t HIDDEN_T = inputOffset + 1;
|
const size_t CELL_T = inputOffset + 2;
|
const size_t HIDDEN_T_GRAD = inputOffset + 3;
|
const size_t CELL_T_GRAD = inputOffset + 4;
|
|
// Extract N
|
const auto N = Input(CELL_T_M_1).size(1);
|
|
// Gates: 1xNxG
|
const auto G = Input(GATES).size(2);
|
const auto D = Input(CELL_T_M_1).size(2);
|
|
CAFFE_ENFORCE_EQ(4 * D, G);
|
const auto* C_prev = Input(CELL_T_M_1).template data<T>();
|
const auto* X = Input(GATES).template data<T>();
|
const auto t = static_cast<OperatorBase*>(this)
|
->Input<Tensor>(TIMESTEP, CPU)
|
.template data<int32_t>()[0];
|
const auto* C = Input(CELL_T).template data<T>();
|
const auto* H = Input(HIDDEN_T).template data<T>();
|
const auto* C_diff = Input(CELL_T_GRAD).template data<T>();
|
const auto* H_diff = Input(HIDDEN_T_GRAD).template data<T>();
|
|
const int32_t* seqLengths = nullptr;
|
if (sequence_lengths_) {
|
CAFFE_ENFORCE_EQ(Input(SEQ_LENGTHS).numel(), N);
|
seqLengths = Input(SEQ_LENGTHS).template data<int32_t>();
|
}
|
|
Output(HIDDEN_T_M_1_GRAD)->ResizeLike(Input(HIDDEN_T_M_1));
|
auto* H_prev_diff = Output(HIDDEN_T_M_1_GRAD)->template mutable_data<T>();
|
Output(CELL_T_M_1_GRAD)->ResizeLike(Input(CELL_T_M_1));
|
auto* C_prev_diff = Output(CELL_T_M_1_GRAD)->template mutable_data<T>();
|
Output(GATES_GRAD)->ResizeLike(Input(GATES));
|
auto* X_diff = Output(GATES_GRAD)->template mutable_data<T>();
|
|
detail::LSTMUnitGradient<T, Context>(
|
N,
|
D,
|
t,
|
C_prev,
|
X,
|
seqLengths,
|
C,
|
H,
|
C_diff,
|
H_diff,
|
drop_states_,
|
H_prev_diff,
|
C_prev_diff,
|
X_diff,
|
forget_bias_,
|
&context_);
|
return true;
|
}
|
|
bool RunOnDevice() override {
|
return DoRunWithType<float>();
|
}
|
|
protected:
|
INPUT_TAGS(HIDDEN_T_M_1, CELL_T_M_1, GATES, SEQ_LENGTHS);
|
// additional input tags are determined dynamically based on whether
|
// sequence_lengths is present.
|
OUTPUT_TAGS(HIDDEN_T_M_1_GRAD, CELL_T_M_1_GRAD, GATES_GRAD);
|
|
float forget_bias_;
|
bool sequence_lengths_;
|
|
private:
|
bool drop_states_;
|
};
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_LSTM_UNIT_OP_H_
|
