wehub-resource-sync
670 lines
22 KiB
C++
670 lines
22 KiB
C++
/*
|
||
* ******************************************************************************
|
||
* *
|
||
* *
|
||
* * This program and the accompanying materials are made available under the
|
||
* * terms of the Apache License, Version 2.0 which is available at
|
||
* * https://www.apache.org/licenses/LICENSE-2.0.
|
||
* *
|
||
* * See the NOTICE file distributed with this work for additional
|
||
* * information regarding copyright ownership.
|
||
* * 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.
|
||
* *
|
||
* * SPDX-License-Identifier: Apache-2.0
|
||
* *****************************************************************************
|
||
*/
|
||
|
||
//
|
||
// @author Yurii Shyrma (iuriish@yahoo.com), created on 15.02.2018, Alex Black
|
||
//
|
||
|
||
#include <system/op_boilerplate.h>
|
||
|
||
|
||
#if NOT_EXCLUDED(OP_gru)
|
||
|
||
// implementation of gated Recurrent Unit cell
|
||
// (cf. https://arxiv.org/abs/1406.1078).
|
||
// Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, Yoshua Bengio
|
||
// "Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation"
|
||
|
||
#include <helpers/MmulHelper.h>
|
||
#include <ops/declarable/CustomOperations.h>
|
||
#include <ops/declarable/helpers/gru.h>
|
||
#include <ops/declarable/helpers/transforms.h>
|
||
|
||
namespace sd {
|
||
namespace ops {
|
||
namespace helpers {
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void gruCell(sd::LaunchContext* context, NDArray* x, NDArray* hI, NDArray* W, NDArray* Wc,
|
||
NDArray* b, NDArray* bc, NDArray* r, NDArray* u, NDArray* c, NDArray* h) {
|
||
// Inputs:
|
||
// x input [bS, nIn], nIn - input size
|
||
// hI previous cell output [bS, nOut], that is at previous time step t-1, nOut - number of units
|
||
// W RU weights - [nIn+nOut, 2*nOut] - reset and update gates
|
||
// Wc C weights - [nIn+nOut, nOut] - cell gate
|
||
// b r and u biases, [2*nOut] - reset and update gates
|
||
// bc c biases, [nOut] - cell gate
|
||
|
||
// Outputs:
|
||
// r Reset gate output [bS, nOut]
|
||
// u Update gate output [bS, nOut]
|
||
// c Cell gate output [bS, nOut]
|
||
// h current cell output [bS, nOut]
|
||
|
||
/***************************************************************************************/
|
||
/************************ THIS IS NOT OPTIMIZED CODE ***********************************/
|
||
/** however it is more math-friendly and convenient for backprop formulas derivation) **/
|
||
|
||
const int bS = x->sizeAt(0);
|
||
const int nIn = x->sizeAt(1);
|
||
const int nOut = hI->sizeAt(1);
|
||
|
||
NDArray *Wrx = (*W)({0, nIn, 0, nOut}); // [nIn, nOut]
|
||
NDArray *Wux = (*W)({0, nIn, nOut, 2 * nOut}); // [nIn, nOut]
|
||
NDArray *Wrh = (*W)({nIn, nIn + nOut, 0, nOut}); // [nOut, nOut]
|
||
NDArray *Wuh = (*W)({nIn, nIn + nOut, nOut, 2 * nOut}); // [nOut, nOut]
|
||
|
||
NDArray *Wcx = (*Wc)({0, nIn, 0, 0}); // reset cell weights [nIn, nOut]
|
||
NDArray *Wch = (*Wc)({nIn, nIn + nOut, 0, 0}); // updates cell weights [nOut, nOut]
|
||
|
||
NDArray *br = (*b)({0, nOut}); // [nOut]
|
||
NDArray *bu = (*b)({nOut, 2 * nOut}); // [nOut]
|
||
|
||
// × means matrix multiplication
|
||
// * means element-wise product or so called Hadamard product
|
||
|
||
// r = sigmoid(x × Wrx + hI × Wrh + br)
|
||
auto xWrx = mmul(*x, *Wrx);
|
||
auto hIWrh = mmul(*hI, *Wrh);
|
||
auto* sum1 = *xWrx + *hIWrh;
|
||
auto* rAssign = (*sum1) + (*br);
|
||
delete sum1;
|
||
delete hIWrh;
|
||
r->assign(rAssign); // [bS, nIn] × [nIn, nOut] + [bS, nOut] × [nOut, nOut] + [nOut] = [bS, nOut]
|
||
delete rAssign;
|
||
r->applyTransform(transform::Sigmoid, r);
|
||
|
||
// u = sigmoid(x × Wux + hI × Wuh + bu)
|
||
auto xWux = mmul(*x, *Wux);
|
||
auto hIWuh = mmul(*hI, *Wuh);
|
||
auto* sum2 = *xWux + *hIWuh;
|
||
auto* uAssign = (*sum2) + (*bu);
|
||
delete sum2;
|
||
delete xWux;
|
||
u->assign(uAssign); // [bS, nIn] × [nIn, nOut] + [bS, nOut] × [nOut, nOut] + [nOut] = [bS, nOut]
|
||
delete uAssign;
|
||
u->applyTransform(transform::Sigmoid, u);
|
||
|
||
// c = tanh(x × Wcx + (r * hI) × Wch + bc)
|
||
auto* rTimesHi = (*r) * (*hI);
|
||
auto xWcx = mmul(*x, *Wcx);
|
||
auto rTimesHiWch = mmul(*rTimesHi, *Wch);
|
||
delete rTimesHi;
|
||
auto* sum3 = *xWcx + *rTimesHiWch;
|
||
auto* cAssign = (*sum3) + (*bc);
|
||
delete sum3;
|
||
delete xWcx;
|
||
delete rTimesHiWch;
|
||
c->assign(cAssign); // [bS, nIn] × [nIn, nOut] + [bS, nOut] × [nOut, nOut] + [nOut] = [bS, nOut]
|
||
delete cAssign;
|
||
c->applyTransform(transform::Tanh, c);
|
||
|
||
// h = (1 - u) * c + u * hI
|
||
auto* uTimesHi = (*u) * (*hI);
|
||
auto* oneMinusU = 1.f - (*u);
|
||
auto* oneMinusUTimesC = (*oneMinusU) * (*c);
|
||
delete oneMinusU;
|
||
auto* hAssign = (*uTimesHi) + (*oneMinusUTimesC);
|
||
delete uTimesHi;
|
||
delete oneMinusUTimesC;
|
||
h->assign(hAssign);
|
||
delete hAssign;
|
||
|
||
delete Wrx;
|
||
delete Wux;
|
||
delete Wrh;
|
||
delete Wuh;
|
||
delete Wcx;
|
||
delete Wch;
|
||
delete br;
|
||
delete bu;
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void gruCell(NDArray* x, NDArray* hI, NDArray* Wx, NDArray* Wh, NDArray* b,
|
||
NDArray* gates, NDArray* h, bool linearBeforeReset) {
|
||
|
||
if(linearBeforeReset) {
|
||
THROW_EXCEPTION("GRU: Linear before reset not implemented. Please set to false.");
|
||
}
|
||
|
||
// Inputs:
|
||
// x input [bS, nIn]
|
||
// hI previous cell output [bS, nOut], that is at previous time step t-1
|
||
// Wx weights for x - [nIn, 3*nOut]
|
||
// Wh weights for h - [nOut, 3*nOut]
|
||
// b biases [3*nOut]
|
||
|
||
// 3*nOut means following sequence: reset, update, cell
|
||
|
||
// Outputs:
|
||
// gates [bS, 3*nOut] = reset gate [bS, nOut] + update gate [bS, nOut] + cell gate [bS, nOut]
|
||
// h current cell output [bS, nOut]
|
||
|
||
// formulas:
|
||
// zr = x × Wxr + hI × Whr + br
|
||
// zu = x × Wxu + hI × Whu + bu
|
||
// r = sigmoid(zr)
|
||
// u = sigmoid(zu)
|
||
// zc = x × Wxc + (r * hI) × Whc + bc
|
||
// c = tanh(zc)
|
||
// h = (1-u)*c + u*hI
|
||
|
||
const int bS = x->sizeAt(0);
|
||
const int nIn = x->sizeAt(1);
|
||
const int nOut = hI->sizeAt(1);
|
||
|
||
NDArray *gatesULike = gates->ulike();
|
||
NDArray temp = *gatesULike;
|
||
MmulHelper::mmul(x, Wx, &temp); // [bS, nIn] × [nIn, 3*nOut] = [bS, 3*nOut]
|
||
temp += *b;
|
||
|
||
MmulHelper::mmul(hI, Wh, gates); // [bS, nOut] × [nOut, 3*nOut] = [bS, 3*nOut]
|
||
|
||
NDArray *ru = (*gates)({0, 0, 0, 2 * nOut}); // [bS, 2*nOut]
|
||
|
||
NDArray *r = (*gates)({0, 0, 0, nOut}); // [bS, nOut]
|
||
NDArray *u = (*gates)({0, 0, nOut, 2 * nOut}); // [bS, nOut]
|
||
NDArray *c = (*gates)({0, 0, 2 * nOut, 3 * nOut}); // [bS, nOut]
|
||
|
||
NDArray *tempView1 = temp({0, 0, 0, 2 * nOut});
|
||
NDArray *tempView2 = temp({0, 0, 2 * nOut, 3 * nOut});
|
||
|
||
// reset and update gates
|
||
*ru += *tempView1;
|
||
ru->applyTransform(transform::Sigmoid, ru);
|
||
|
||
// cell gate
|
||
auto* cTimesR = (*c) * (*r);
|
||
auto* cAssign = (*cTimesR) + (*tempView2);
|
||
delete cTimesR;
|
||
c->assign(cAssign);
|
||
delete cAssign;
|
||
c->applyTransform(transform::Tanh, c);
|
||
|
||
// h = (1-u)*c + u*hI
|
||
auto* uTimesHi = (*u) * (*hI);
|
||
auto* oneMinusU = 1.f - (*u);
|
||
auto* oneMinusUTimesC = (*oneMinusU) * (*c);
|
||
delete oneMinusU;
|
||
auto* hAssign = (*uTimesHi) + (*oneMinusUTimesC);
|
||
delete uTimesHi;
|
||
delete oneMinusUTimesC;
|
||
h->assign(hAssign);
|
||
delete hAssign;
|
||
|
||
delete gatesULike;
|
||
delete ru;
|
||
delete r;
|
||
delete u;
|
||
delete c;
|
||
delete tempView1;
|
||
delete tempView2;
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void gruTimeLoop(sd::LaunchContext* context, NDArray* x, NDArray* hI, NDArray* Wx, NDArray* Wh,
|
||
NDArray* b, NDArray* h, bool linearBeforeReset) {
|
||
// sL means time steps
|
||
|
||
// x input [sL, bS, nIn]
|
||
// hI initial cell output (at time step = 0) [bS, nOut]
|
||
// Wx input-to-hidden weights, [nIn, 3*nOut]
|
||
// Wh hidden-to-hidden weights, [nOut, 3*nOut]
|
||
// b biases, [3*nOut]
|
||
|
||
// h cell outputs at each time step [sL, bS, nOut]
|
||
|
||
const int sL = x->sizeAt(0);
|
||
const int bS = x->sizeAt(1);
|
||
const int nOut = hI->sizeAt(1);
|
||
|
||
std::vector<LongType> shape = {bS, 3 * nOut};
|
||
NDArray gates(h->ordering(), shape, h->dataType(), context);
|
||
|
||
auto xSet = x->allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nIn]
|
||
auto hSet = h->allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nOut]
|
||
|
||
// time loop
|
||
for (int t = 0; t < sL; ++t) {
|
||
gruCell(xSet.at(t), t == 0 ? hI : hSet.at(t - 1), Wx, Wh, b, &gates, hSet.at(t), linearBeforeReset);
|
||
}
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void gruCellBp(sd::LaunchContext* context, NDArray* x, NDArray* hLast, NDArray* W, NDArray* Wc,
|
||
NDArray* b, NDArray* bc, NDArray* dLdr, NDArray* dLdu, NDArray* dLdc,
|
||
NDArray* dLdh, NDArray* dLdx, NDArray* dLdhLast, NDArray* dLdW, NDArray* dLdWc, NDArray* dLdb,
|
||
NDArray* dLdbc) {
|
||
// Inputs:
|
||
// x input [bS, iS]
|
||
// hLast previous cell output [bS, nU], that is at previous time step t-1
|
||
// W weights - [iS+nU, 2*nU] - reset and update gates
|
||
// Wc C weights - [iS+nU, nU] - cell gate
|
||
// b r and u biases, [2*nU] - reset and update gates
|
||
// bc c biases, [nU] - cell gate
|
||
// dLdr gradient wrt reset gate, [bS, nU]
|
||
// dLdu gradient wrt update gate, [bS, nU]
|
||
// dLdc gradient wrt cell state, [bS, nU]
|
||
// dLdh gradient wrt current cell output, [bS, nU]
|
||
|
||
// Outputs:
|
||
// dLdx gradient wrt x, [bS, iS],
|
||
// dLdhLast gradient wrt hLast, [bS, nU]
|
||
// dLdW gradient wrt W, [iS+nU, 2*nU]
|
||
// dLdWc gradient wrt Wc, [iS+nU, nU]
|
||
// dLdb gradient wrt bru [2*nU]
|
||
// dLdbc gradient wrt bc [nU]
|
||
|
||
// * means element-wise product or so called Hadamard product
|
||
// × means matrix multiplication
|
||
|
||
/************************************************************************************************/
|
||
/******************************* THIS IS NOT OPTIMIZED CODE *************************************/
|
||
/*** aim is to have math-readable code in order to keep track of backprop formulas derivation ***/
|
||
|
||
const int bS = x->sizeAt(0);
|
||
const int iS = x->sizeAt(1);
|
||
const int nU = hLast->sizeAt(1);
|
||
|
||
NDArray *xT = x->transpose(); // [iS, bS]
|
||
NDArray *hLastT = hLast->transpose(); // [nU, bS]
|
||
|
||
NDArray *Wrx = (*W)({0, iS, 0, nU}); // [iS, nU]
|
||
NDArray *Wux = (*W)({0, iS, nU, 2 * nU}); // [iS, nU]
|
||
NDArray *Wrh = (*W)({iS, iS + nU, 0, nU}); // [nU, nU]
|
||
NDArray *Wuh = (*W)({iS, iS + nU, nU, 2 * nU}); // [nU, nU]
|
||
|
||
NDArray *Wcx = (*Wc)({0, iS, 0, 0}); // reset cell weights [iS, nU]
|
||
NDArray *Wch = (*Wc)({iS, iS + nU, 0, 0}); // updates cell weights [nU, nU]
|
||
|
||
NDArray *br = (*b)({0, nU}); // [nU]
|
||
NDArray *bu = (*b)({nU, 2 * nU}); // [nU]
|
||
|
||
NDArray *WrxT = Wrx->transpose(); // [nU, iS]
|
||
NDArray *WuxT = Wux->transpose(); // [nU, iS]
|
||
NDArray *WrhT = Wrh->transpose(); // [nU, nU]
|
||
NDArray *WuhT = Wuh->transpose(); // [nU, nU]
|
||
|
||
NDArray *WcxT = Wcx->transpose(); // [nU, iS]
|
||
NDArray *WchT = Wch->transpose(); // [nU, nU]
|
||
|
||
NDArray *dLdWrx = (*dLdW)({0, iS, 0, nU}); // [iS, nU]
|
||
NDArray *dLdWux = (*dLdW)({0, iS, nU, 2 * nU}); // [iS, nU]
|
||
NDArray *dLdWrh = (*dLdW)({iS, iS + nU, 0, nU}); // [nU, nU]
|
||
NDArray *dLdWuh = (*dLdW)({iS, iS + nU, nU, 2 * nU}); // [nU, nU]
|
||
|
||
NDArray *dLdWcx = (*dLdWc)({0, iS, 0, 0}); // [iS, nU]
|
||
NDArray *dLdWch = (*dLdWc)({iS, iS + nU, 0, 0}); // [nU, nU]
|
||
|
||
NDArray *dLdbr = (*dLdb)({0, nU}); // [nU]
|
||
NDArray *dLdbu = (*dLdb)({nU, 2 * nU}); // [nU]
|
||
|
||
// ***** feed forward step ***** //
|
||
|
||
// r = sigmoid(x × Wrx + hLast × Wrh + br)
|
||
auto xWrx = mmul(*x, *Wrx);
|
||
auto hLastWrh = mmul(*hLast, *Wrh);
|
||
auto* sum1 = *xWrx + *hLastWrh;
|
||
auto* rTemp = (*sum1) + (*br);
|
||
delete sum1;
|
||
NDArray r = *rTemp; // [bS, iS] × [iS, nU] + [bS, nU] × [nU, nU] + [nU] = [bS, nU]
|
||
delete rTemp;
|
||
r.applyTransform(transform::Sigmoid, &r);
|
||
|
||
// u = sigmoid(x × Wux + hLast × Wuh + bu)
|
||
auto xWux = mmul(*x, *Wux);
|
||
auto hLastWuh = mmul(*hLast, *Wuh);
|
||
auto* sum2 = *xWux + *hLastWuh;
|
||
auto* uTemp = (*sum2) + (*bu);
|
||
delete sum2;
|
||
NDArray u = *uTemp; // [bS, iS] × [iS, nU] + [bS, nU] × [nU, nU] + [nU] = [bS, nU]
|
||
delete uTemp;
|
||
delete xWux;
|
||
u.applyTransform(transform::Sigmoid, &u);
|
||
|
||
// c = tanh(x × Wcx + (r * hLast) × Wch + bc)
|
||
auto* rTimesHLast2 = r * (*hLast);
|
||
auto xWcx = mmul(*x, *Wcx);
|
||
auto rTimesHLast2Wch = mmul(*rTimesHLast2, *Wch);
|
||
delete rTimesHLast2;
|
||
auto* sum3 = *xWcx + *rTimesHLast2Wch;
|
||
auto* cTemp = (*sum3) + (*bc);
|
||
delete sum3;
|
||
delete xWcx;
|
||
delete rTimesHLast2Wch;
|
||
NDArray c = *cTemp; // [bS, iS] × [iS, nU] + [bS, nU] × [nU, nU] + [nU] = [bS, nU]
|
||
delete cTemp;
|
||
c.applyTransform(transform::Tanh, &c);
|
||
|
||
// h = (1 - u) * c + u * hPrev
|
||
|
||
// ***** back prop step ***** //
|
||
|
||
auto* hLastMinusC = (*hLast) - c;
|
||
auto* oneMinusU = 1.f - u;
|
||
auto* dudZu = u * (*oneMinusU);
|
||
delete oneMinusU;
|
||
auto* oneMinusR = 1.f - r;
|
||
auto* drdZr = r * (*oneMinusR);
|
||
delete oneMinusR;
|
||
auto* cSquared = c * c;
|
||
auto* oneMinusCSquared = 1.f - (*cSquared);
|
||
delete cSquared;
|
||
auto dcdZc = *oneMinusCSquared;
|
||
delete oneMinusCSquared;
|
||
auto* dLdZc = (*dLdc) * dcdZc;
|
||
auto* dLdZu = (*dLdu) * (*dudZu);
|
||
delete dudZu;
|
||
auto* dLdZr = (*dLdr) * (*drdZr);
|
||
delete drdZr;
|
||
|
||
NDArray *dhdc = 1.f - u; // [bS, nU]
|
||
NDArray dhdu = *hLastMinusC; // [bS, nU]
|
||
delete hLastMinusC;
|
||
|
||
// dLdx = dLdZu × WuxT + dLdZc × WcxT + dLdZr × WrxT
|
||
auto dLdZuWuxT = mmul(*dLdZu, *WuxT);
|
||
auto dLdZcWcxT = mmul(*dLdZc, *WcxT);
|
||
auto dLdZrWrxT = mmul(*dLdZr, *WrxT);
|
||
auto* temp1 = *dLdZuWuxT + *dLdZcWcxT;
|
||
auto* dLdxTemp = (*temp1) + *dLdZrWrxT;
|
||
delete temp1;
|
||
|
||
delete dLdZuWuxT;
|
||
delete dLdZcWcxT;
|
||
delete dLdZrWrxT;
|
||
dLdx->assign(dLdxTemp); // [bS, iS]
|
||
delete dLdxTemp;
|
||
|
||
// dldZTimeR = dLdZc * r
|
||
auto* dldZTimeR = (*dLdZc) * r;
|
||
|
||
// dLdhLast = dLdh * u + dLdZu × WuhT + dldZTimeR × WchT + dLdZr × WrhT
|
||
auto* dLdhTimesU = (*dLdh) * u;
|
||
auto dLdZuWuhT = mmul(*dLdZu, *WuhT);
|
||
auto dldZTimeRWchT = mmul(*dldZTimeR, *WchT);
|
||
auto dLdZrWrhT = mmul(*dLdZr, *WrhT);
|
||
auto* temp2 = (*dLdhTimesU) + *dLdZuWuhT;
|
||
delete dLdhTimesU;
|
||
delete dLdZuWuhT;
|
||
auto* temp3 = (*temp2) + *dldZTimeRWchT;
|
||
delete temp2;
|
||
delete dldZTimeRWchT;
|
||
auto* dLdhLastTemp = (*temp3) + *dLdZrWrhT;
|
||
delete temp3;
|
||
delete dLdZrWrhT;
|
||
dLdhLast->assign(dLdhLastTemp); // [bS, nU]
|
||
delete dLdhLastTemp;
|
||
|
||
// dLdWrx = xT × dLdZr
|
||
auto dLdWrxTemp = mmul(*xT, *dLdZr);
|
||
dLdWrx->assign(dLdWrxTemp); // [iS, bS] × [bS, nU] = [iS, nU]
|
||
delete dLdWrxTemp;
|
||
// dLdWrh = hLastT × dLdZr
|
||
auto dLdWrhTemp = mmul(*hLastT, *dLdZr);
|
||
dLdWrh->assign(dLdWrhTemp); // [nU, bS] × [bS, nU] = [nU, nU]
|
||
delete dLdWrhTemp;
|
||
// dLdWux = xT × dLdZu
|
||
auto dLdWuxTemp = mmul(*xT, *dLdZu);
|
||
dLdWux->assign(dLdWuxTemp); // [iS, bS] × [bS, nU] = [iS, nU]
|
||
delete dLdWuxTemp;
|
||
// dLdWuh = hLastT × dLdZu
|
||
auto dLdWuhTemp = mmul(*hLastT, *dLdZu);
|
||
dLdWuh->assign(dLdWuhTemp); // [nU, bS] × [bS, nU] = [nU, nU]
|
||
delete dLdWuhTemp;
|
||
// dLdWcx = xT × dLdZc
|
||
auto dLdWcxTemp = mmul(*xT, *dLdZc);
|
||
dLdWcx->assign(dLdWcxTemp); // [iS, bS] × [bS, nU] = [iS, nU]
|
||
delete dLdWcxTemp;
|
||
// dLdWch = (r * hLast)T × dLdZc
|
||
auto* rTimesHLast = r * (*hLast);
|
||
NDArray* rTimesHLastT = rTimesHLast->transpose();
|
||
delete rTimesHLast;
|
||
auto dLdWchTemp = mmul(*rTimesHLastT, *dLdZc);
|
||
dLdWch->assign(dLdWchTemp); // [nU, bS] × [bS, nU] = [nU, nU]
|
||
delete dLdWchTemp;
|
||
// Calculate reduction for bias gradients
|
||
std::vector<sd::LongType> zeroVec = {0};
|
||
auto* dLdbrTemp = dLdZr->reduceAlongDimension(reduce::Sum, &zeroVec);
|
||
dLdbr->assign(dLdbrTemp); // [nU]
|
||
delete dLdbrTemp;
|
||
|
||
auto* dLdbuTemp = dLdZu->reduceAlongDimension(reduce::Sum, &zeroVec);
|
||
dLdbu->assign(dLdbuTemp); // [nU]
|
||
delete dLdbuTemp;
|
||
|
||
auto* dLdbcTemp = dLdZc->reduceAlongDimension(reduce::Sum, &zeroVec);
|
||
dLdbc->assign(dLdbcTemp); // [nU]
|
||
delete dLdbcTemp;
|
||
|
||
delete dhdc;
|
||
delete dLdZc;
|
||
delete dLdZu;
|
||
delete dLdZr;
|
||
delete dldZTimeR;
|
||
delete Wrx;
|
||
delete Wux;
|
||
delete Wrh;
|
||
delete Wuh;
|
||
delete Wcx;
|
||
delete Wch;
|
||
delete br;
|
||
delete bu;
|
||
delete WrxT;
|
||
delete WuxT;
|
||
delete WrhT;
|
||
delete WuhT;
|
||
delete WcxT;
|
||
delete WchT;
|
||
delete dLdWrx;
|
||
delete dLdWux;
|
||
delete dLdWrh;
|
||
delete dLdWuh;
|
||
delete dLdWcx;
|
||
delete dLdWch;
|
||
delete dLdbr;
|
||
delete dLdbu;
|
||
delete xT;
|
||
delete hLastT;
|
||
delete rTimesHLastT;
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void gruCellBp(sd::LaunchContext* context, NDArray* x, NDArray* hI, NDArray* Wx, NDArray* Wh,
|
||
NDArray* b, NDArray* dLdh, NDArray* gates, NDArray* dLdx, NDArray* dLdhI,
|
||
NDArray* dLdWx, NDArray* dLdWh, NDArray* dLdb) {
|
||
// Inputs:
|
||
// x input [bS, nIn]
|
||
// hI previous cell output [bS, nOut], that nIn at previous time step t-1
|
||
// Wx input-to-hidden weights - [nIn, 3*nOut]
|
||
// Wh hidden-to-hidden weights - [nOut, 3*nOut]
|
||
// b biases, [3*nOut] - reset and update gates
|
||
// dLdh gradient vs. ff output, [bS, nOut]
|
||
|
||
// Outputs:
|
||
// dLdx gradient vs. x, [bS, nIn],
|
||
// dLdhI gradient vs. hI, [bS, nOut]
|
||
// dLdWx gradient vs. W, [nIn, 3*nOut]
|
||
// dLdWh gradient vs. Wc, [nOut, 3*nOut]
|
||
// dLdb gradient vs. b [3*nOut]
|
||
|
||
// 3*nOut means following sequence: reset, update, cell
|
||
|
||
// * means element-wise product or so called Hadamard product
|
||
// × means matrix multiplication
|
||
|
||
const int nOut = hI->sizeAt(1);
|
||
|
||
NDArray *gatesULike = gates->ulike();
|
||
NDArray dLdz = *gatesULike; // [bS, 3*nOut]
|
||
|
||
NDArray *dLdzru = dLdz({0, 0, 0, 2 * nOut}); // [bS, 2*nOut]
|
||
|
||
NDArray *dLdzr = dLdz({0, 0, 0, nOut}); // [bS, nOut]
|
||
NDArray *dLdzu = dLdz({0, 0, nOut, 2 * nOut}); // [bS, nOut]
|
||
NDArray *dLdzc = dLdz({0, 0, 2 * nOut, 3 * nOut}); // [bS, nOut]
|
||
|
||
NDArray *r = (*gates)({0, 0, 0, nOut}); // [bS, nOut]
|
||
NDArray *u = (*gates)({0, 0, nOut, 2 * nOut}); // [bS, nOut]
|
||
NDArray *c = (*gates)({0, 0, 2 * nOut, 3 * nOut}); // [bS, nOut]
|
||
|
||
NDArray *WhView = (*Wh)({0, 0, 2 * nOut, 3 * nOut});
|
||
NDArray *WhcT = WhView->transpose();
|
||
|
||
if (dLdh) *dLdhI += *dLdh;
|
||
|
||
auto* oneMinusU = 1 - (*u); // [bS, nOut]
|
||
|
||
// dLdzc = dLdhI * (1-u) * (1-c²)
|
||
auto* cSquared = (*c) * (*c);
|
||
auto* oneMinusCSquared = 1.f - (*cSquared);
|
||
delete cSquared;
|
||
auto* temp1 = (*dLdhI) * (*oneMinusU);
|
||
auto* dLdzcTemp = (*temp1) * (*oneMinusCSquared);
|
||
delete temp1;
|
||
delete oneMinusCSquared;
|
||
dLdzc->assign(dLdzcTemp); // [bS, nOut]
|
||
delete dLdzcTemp;
|
||
|
||
// dLdzu = dLdhI * (hI - c) * u * (1-u)
|
||
auto* hIMinusC = (*hI) - (*c);
|
||
auto* uTimesOneMinusU = (*u) * (*oneMinusU);
|
||
auto* temp2 = (*dLdhI) * (*hIMinusC);
|
||
auto* dLdzuTemp = (*temp2) * (*uTimesOneMinusU);
|
||
delete temp2;
|
||
delete hIMinusC;
|
||
delete uTimesOneMinusU;
|
||
dLdzu->assign(dLdzuTemp); // [bS, nOut]
|
||
delete dLdzuTemp;
|
||
delete oneMinusU;
|
||
|
||
// dLdzr = (dLdzc * hI * r * (1-r)) × WhcT
|
||
auto* oneMinusR = 1 - (*r);
|
||
auto* rTimesOneMinusR = (*r) * (*oneMinusR);
|
||
delete oneMinusR;
|
||
auto* temp3 = (*dLdzc) * (*hI);
|
||
auto* temp4 = (*temp3) * (*rTimesOneMinusR);
|
||
delete temp3;
|
||
delete rTimesOneMinusR;
|
||
MmulHelper::mmul(temp4, WhcT, dLdzr); // [bS, nOut] x [nOut, nOut] = [bS, nOut]
|
||
delete temp4;
|
||
|
||
// dLdx = dLdz × WxT
|
||
NDArray *WxT = Wx->transpose();
|
||
MmulHelper::mmul(&dLdz, WxT, dLdx); // [bS, 3*nOut] x [3*nOut, nIn] = [bS, nIn]
|
||
|
||
// dLdWx += xT × dLdz
|
||
NDArray *xT = x->transpose();
|
||
auto dLdWxAdd = mmul(*xT, dLdz);
|
||
*dLdWx += *dLdWxAdd; // [nIn, bS] x [bS, 3*nOut] = [nIn, 3*nOut]
|
||
|
||
delete dLdWxAdd;
|
||
// dLdb += sum(dLdz, axis=0)
|
||
std::vector<sd::LongType> zeroVec = {0};
|
||
auto* dLdbAdd = dLdz.reduceAlongDimension(reduce::Sum, &zeroVec);
|
||
*dLdb += (*dLdbAdd); // [bS, 3*nOut] -> reduce -> [3*nOut];
|
||
delete dLdbAdd;
|
||
|
||
*dLdzc *= (*r);
|
||
|
||
// dLdhI = dLdhI * u + dLdz × WhT
|
||
NDArray *WhT = Wh->transpose();
|
||
auto* dLdhIU = (*dLdhI) * (*u);
|
||
auto mmulResult = mmul(dLdz, *WhT);
|
||
auto* dLdhIAssign = (*dLdhIU) + *mmulResult;
|
||
delete dLdhIU;
|
||
delete mmulResult;
|
||
dLdhI->assign(dLdhIAssign); // [bS, 3*nOut] x [3*nOut, nOut] = [bS, nOut]
|
||
delete dLdhIAssign;
|
||
|
||
// dLdWh += hIT × dLdz
|
||
NDArray *hITranspose = hI->transpose();
|
||
auto dLdWhAdd = mmul(*hITranspose, dLdz);
|
||
*dLdWh += *dLdWhAdd; // [nOut, bS] x [bS, 3*nOut] = [nOut, 3*nOut]
|
||
delete dLdWhAdd;
|
||
|
||
delete gatesULike;
|
||
delete dLdzru;
|
||
delete dLdzr;
|
||
delete dLdzu;
|
||
delete dLdzc;
|
||
delete r;
|
||
delete u;
|
||
delete c;
|
||
delete WhView;
|
||
delete WhcT;
|
||
delete hITranspose;
|
||
delete WhT;
|
||
delete xT;
|
||
delete WxT;
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void gruTimeLoopBp(sd::LaunchContext* context, NDArray* x, NDArray* hI, NDArray* Wx,
|
||
NDArray* Wh, NDArray* b, NDArray* dLdh, NDArray* dLdx, NDArray* dLdhI,
|
||
NDArray* dLdWx, NDArray* dLdWh, NDArray* dLdb) {
|
||
// sL means time steps
|
||
|
||
// x input [sL, bS, nIn]
|
||
// hI initial cell output (at time step = 0) [bS, nOut]
|
||
// Wx input-to-hidden weights, [nIn, 3*nOut]
|
||
// Wh hidden-to-hidden weights, [nOut, 3*nOut]
|
||
// b biases, [3*nOut]
|
||
// dLdh gradient vs. ff output, [sL, bS, nOut]
|
||
|
||
// dLdx gradient vs. x, [sL, bS, nIn],
|
||
// dLdhI gradient vs. hI, [bS, nOut]
|
||
// dLdWx gradient vs. W, [nIn, 3*nOut]
|
||
// dLdWh gradient vs. Wc, [nOut, 3*nOut]
|
||
// dLdb gradient vs. b [3*nOut]
|
||
|
||
const int sL = x->sizeAt(0);
|
||
const int bS = x->sizeAt(1);
|
||
const int nOut = hI->sizeAt(1);
|
||
|
||
std::vector<sd::LongType> shape = {bS, 3 * nOut};
|
||
std::vector<sd::LongType> hShape = {sL + 1, bS, nOut};
|
||
NDArray gates(x->ordering(), shape, dLdh->dataType(), x->getContext());
|
||
NDArray h(x->ordering(), hShape, dLdh->dataType(), x->getContext());
|
||
|
||
auto xSet = x->allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nIn]
|
||
auto dLdhSet = dLdh->allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nOut]
|
||
auto hSet = h.allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nOut]
|
||
auto gatesSet = gates.allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nOut]
|
||
auto dLdxSet = dLdx->allTensorsAlongDimension({1, 2}); // sub-arrays with shape [bS, nIn]
|
||
|
||
hSet.at(0)->assign(hI);
|
||
|
||
// forward time loop
|
||
for (int t = 0; t < sL; ++t) gruCell(xSet.at(t), hSet.at(t), Wx, Wh, b, gatesSet.at(t), hSet.at(t + 1), false);
|
||
|
||
// backward time loop
|
||
for (int t = sL - 1; t >= 0; --t)
|
||
gruCellBp(context, xSet.at(t), hSet.at(t), Wx, Wh, b, dLdhSet.at(t), gatesSet.at(t), dLdxSet.at(t), dLdhI, dLdWx,
|
||
dLdWh, dLdb);
|
||
}
|
||
|
||
} // namespace helpers
|
||
} // namespace ops
|
||
} // namespace sd
|
||
|
||
#endif
|