提交 f10710b0 编写于 作者: T tensor-tang

move seq peephole if out of loop

上级 2f3b4989
...@@ -272,6 +272,10 @@ class FuisonLSTMKernel : public framework::OpKernel<T> { ...@@ -272,6 +272,10 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
T* hidden_out_data = hidden_out->mutable_data<T>(place); T* hidden_out_data = hidden_out->mutable_data<T>(place);
T* cell_out_data = cell_out->mutable_data<T>(place); T* cell_out_data = cell_out->mutable_data<T>(place);
auto blas = math::GetBlas<DeviceContext, T>(ctx);
math::FCCompute<DeviceContext, T>(blas, total_T, D4, M, x_data, wx_data,
xx_data, bias->data<T>());
// for peephole only
Tensor checked_cell; Tensor checked_cell;
T* checked_cell_data = nullptr; T* checked_cell_data = nullptr;
if (use_peepholes) { if (use_peepholes) {
...@@ -279,9 +283,6 @@ class FuisonLSTMKernel : public framework::OpKernel<T> { ...@@ -279,9 +283,6 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
checked_cell_data = checked_cell.mutable_data<T>({2, D}, place); checked_cell_data = checked_cell.mutable_data<T>({2, D}, place);
} }
auto blas = math::GetBlas<DeviceContext, T>(ctx);
math::FCCompute<DeviceContext, T>(blas, total_T, D4, M, x_data, wx_data,
xx_data, bias->data<T>());
int xx_offset = D4; int xx_offset = D4;
int gate_offset = D; int gate_offset = D;
if (is_reverse) { if (is_reverse) {
...@@ -299,6 +300,26 @@ class FuisonLSTMKernel : public framework::OpKernel<T> { ...@@ -299,6 +300,26 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
cell_out_data = cell_out_data + gate_offset; cell_out_data = cell_out_data + gate_offset;
}; };
#define GEMM_WH_ADDON \
blas.GEMM(CblasNoTrans, CblasNoTrans, 1, D4, D, static_cast<T>(1), \
prev_h_data, D, wh_data, D4, static_cast<T>(1), xx_data, D4)
#define GET_Ct \
/* C_t = C_t-1 * fgated + cand_gated * igated*/ \
act_cand(D, xx_data, xx_data); \
blas.VMUL(D, xx_data, xx_data + D, xx_data + D); \
blas.VMUL(D, prev_c_data, xx_data + D2, xx_data + D2); \
blas.VADD(D, xx_data + D, xx_data + D2, cell_out_data)
#define GET_Ht_AND_MOVE \
/* H_t = act_cell(C_t) * ogated */ \
act_cell(D, cell_out_data, xx_data + D2); \
blas.VMUL(D, xx_data + D2, xx_data + D3, hidden_out_data); \
/* get prev and move*/ \
prev_h_data = hidden_out_data; \
prev_c_data = cell_out_data; \
move_step()
for (int i = 0; i < N; ++i) { for (int i = 0; i < N; ++i) {
int bid = is_reverse ? N - 1 - i : i; int bid = is_reverse ? N - 1 - i : i;
int seq_len = x_lod[0][bid + 1] - x_lod[0][bid]; int seq_len = x_lod[0][bid + 1] - x_lod[0][bid];
...@@ -312,67 +333,49 @@ class FuisonLSTMKernel : public framework::OpKernel<T> { ...@@ -312,67 +333,49 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
// W_ch, W_ih, W_fh, W_oh // W_ch, W_ih, W_fh, W_oh
act_gate(D, xx_data + D, xx_data + D); act_gate(D, xx_data + D, xx_data + D);
act_cand(D, xx_data, xx_data); act_cand(D, xx_data, xx_data);
// C_t = input * tilde // C_t = igated * cgated
blas.VMUL(D, xx_data, xx_data + D, cell_out_data); blas.VMUL(D, xx_data, xx_data + D, cell_out_data);
// H_t = act_state(cellout) * outgate // get outgated
if (use_peepholes) { if (use_peepholes) {
// + W_oc * C_t for peephole connection // put W_oc * C_t on igated
// put result on W_ih
blas.VMUL(D, wc_data + D2, cell_out_data, xx_data + D); blas.VMUL(D, wc_data + D2, cell_out_data, xx_data + D);
blas.VADD(D, xx_data + D, xx_data + D3, xx_data + D3); blas.VADD(D, xx_data + D, xx_data + D3, xx_data + D3);
} }
act_gate(D, xx_data + D3, xx_data + D3); act_gate(D, xx_data + D3, xx_data + D3);
act_cell(D, cell_out_data, xx_data + D2); GET_Ht_AND_MOVE;
blas.VMUL(D, xx_data + D2, xx_data + D3, hidden_out_data);
// prev
prev_h_data = hidden_out_data;
prev_c_data = cell_out_data;
tstart = 1; tstart = 1;
move_step();
} }
for (int step = tstart; step < seq_len; ++step) { if (use_peepholes) {
// + W_h * H_t-1 for (int step = tstart; step < seq_len; ++step) {
blas.GEMM(CblasNoTrans, CblasNoTrans, 1, D4, D, static_cast<T>(1), GEMM_WH_ADDON;
prev_h_data, D, wh_data, D4, static_cast<T>(1), xx_data, D4); // get fgated and igated
// W_ch, W_ih, W_fh, W_oh
if (use_peepholes) {
// + W_ic|W_fc * C_t-1 for peephole connection
blas.VMUL(D, wc_data, prev_c_data, checked_cell_data); blas.VMUL(D, wc_data, prev_c_data, checked_cell_data);
blas.VMUL(D, wc_data + D, prev_c_data, checked_cell_data + D); blas.VMUL(D, wc_data + D, prev_c_data, checked_cell_data + D);
blas.VADD(D2, checked_cell_data, xx_data + D, xx_data + D); blas.VADD(D2, checked_cell_data, xx_data + D, xx_data + D);
act_gate(D2, xx_data + D, xx_data + D); act_gate(D2, xx_data + D, xx_data + D);
} else { GET_Ct;
act_gate(D3, xx_data + D, xx_data + D);
}
// a = I_t * act_cand(ch)
act_cand(D, xx_data, xx_data);
blas.VMUL(D, xx_data, xx_data + D, xx_data + D);
// b = C_t-1 * F_t
blas.VMUL(D, prev_c_data, xx_data + D2, xx_data + D2);
// C_t = a + b
blas.VADD(D, xx_data + D, xx_data + D2, cell_out_data);
// H_t = act_cell(C_t) * act_gate(O_c += C_t * W_oc) // get ogated
if (use_peepholes) {
// put result on W_ih
blas.VMUL(D, wc_data + D2, cell_out_data, xx_data + D); blas.VMUL(D, wc_data + D2, cell_out_data, xx_data + D);
blas.VADD(D, xx_data + D, xx_data + D3, xx_data + D3); blas.VADD(D, xx_data + D, xx_data + D3, xx_data + D3);
act_gate(D, xx_data + D3, xx_data + D3); act_gate(D, xx_data + D3, xx_data + D3);
} GET_Ht_AND_MOVE;
act_cell(D, cell_out_data, xx_data + D2); } // for seqlen
blas.VMUL(D, xx_data + D2, xx_data + D3, hidden_out_data); } else {
for (int step = tstart; step < seq_len; ++step) {
// prev GEMM_WH_ADDON;
prev_h_data = hidden_out_data; // W_ch, W_ih, W_fh, W_oh
prev_c_data = cell_out_data; act_gate(D3, xx_data + D, xx_data + D);
GET_Ct;
move_step(); GET_Ht_AND_MOVE;
} // for seqlen } // for seqlen
} // for batch }
} // for batch
#undef GET_Ht_AND_MOVE
#undef GEMM_WH_ADDON
#undef GET_Ct
} }
void BatchCompute(const framework::ExecutionContext& ctx) const { void BatchCompute(const framework::ExecutionContext& ctx) const {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册