/******************************************************************************* * Copyright (c) 2015-2018 Skymind, Inc. * * 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. * * 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 // // implementation of gated Recurrent Unit cell // (cf. http://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 #include #include #include namespace nd4j { namespace ops { namespace helpers { ////////////////////////////////////////////////////////////////////////// void gruCell(nd4j::LaunchContext * context, const NDArray* x, const NDArray* hLast, const NDArray* Wru, const NDArray* Wc, const NDArray* bru, const NDArray* bc, NDArray* r, NDArray* u, NDArray* c, NDArray* h) { //Inputs: // x input [bS, nIn], nIn - input size // hLast previous cell output [bS, nUn], that is at previous time step t-1, nUn - number of units // Wru RU weights - [nIn+nUn, 2*nUn] - reset and update gates // Wc C weights - [nIn+nUn, nUn] - cell gate // bru r and u biases, [2*nUn] - reset and update gates // bc c biases, [nUn] - cell gate //Outputs: // r Reset gate output [bS, nUn] // u Update gate output [bS, nUn] // c Cell gate output [bS, nUn] // h current cell output [bS, nUn] /***************************************************************************************/ /************************ THIS IS NOT OPTIMAZED 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 nUn = hLast->sizeAt(1); NDArray Wr = (*Wru)({0,nIn, 0,0}); // reset gates weights [nIn, 2*nUn] NDArray Wu = (*Wru)({nIn,nIn+nUn, 0,0}); // updates gates weights [nUn, 2*nUn] NDArray Wcr = (*Wc)({0,nIn, 0,0}); // reset cell weights [nIn, nUn] NDArray Wcu = (*Wc)({nIn,nIn+nUn, 0,0}); // updates cell weights [nUn, nUn] // gates = sigmoid(x*Wr + hLast*Wu + br + bu) NDArray gates = mmul(*x, Wr) + mmul(*hLast, Wu) + *bru; // [bS, nIn] * [nIn, 2*nUn] + [bS, nUn] * [nUn, 2*nUn] + [2*nUn] = [bS, 2*nUn] gates.applyTransform(transform::Sigmoid); // reset gate r->assign(gates({0,0, 0,nUn})); // [bS, nUn] // update gate u->assign(gates({0,0, nUn,2*nUn})); // [bS, nUn] // cell gate c = activation(x*Wcr + (r◦hlast)*Wcu + bc) c->assign(mmul(*x, Wcr) + mmul(*r * *hLast, Wcu) + *bc); // [bS, nIn] * [nIn, nUn] + [bS, nUn] * [nUn, nUn] + [nUn] = [bS, nUn] c->applyTransform(transform::Tanh); // cell output h->assign(*u * *hLast + (1.f - *u) * *c); /***************************************************************************************/ /********************** THIS MORE OPTIMAZED CODE (except concat ) **********************/ /***************************************************************************************/ /* //Concat inputs: x + hLast : [bs, nIn + nUn] NDArray xhConcat(x->ordering(), {bS, nIn + nUn}, x->dataType(), context); // concat([bs, nIn], [bs, nUn]) -> [bs, nIn + nUn] helpers::concat(context, {const_cast(x), const_cast(hLast)}, xhConcat, {1}); //mmul for reset and update gates: (x * weight_ux + hLast * weight_xr + b_u) auto m = mmul(xhConcat, *Wru) + *bru ; // [bs, nIn+nUn] * [nIn+nUn, 2*nUn] = [bs, 2*nUn] // m += *bru; sigmoidInplace(m); //sigmoid(rz) and sigmoid(uz) r->assign(m({0,0, 0, nUn})); u->assign(m({0,0, nUn, 2*nUn})); // hLast = hLast * r xhConcat({0,0, nIn, nIn+nUn}) *= *r; //c = tanh(x * weight_cx + (hLast .* r) * weight_cr + b_c) MmulHelper::mmul(&xhConcat, Wc, c, 1.0, 0.0); //c = 1.0 * xhConcat * Wc + 0.0 * c *c += *bc; tanhInplace(*c); //Output: h = (1-u).*c + u .* hPrev //auto hResult = (*u) * (*hLast) + (1.0f - *u) * (*c); const_cast(h)->assign(&hResult); u->applyPairwiseTransform(pairwise::Multiply, hLast, h, nullptr); //h = u * hLast auto temp = (1.0f - *u); temp *= (*c); (*h) += temp; */ } ////////////////////////////////////////////////////////////////////////// void gruTimeLoop(nd4j::LaunchContext * context, const NDArray* x, const NDArray* h0, const NDArray* Wx, const NDArray* Wh, const NDArray* b, NDArray* h) { // x input [time, bS, iS] // h0 initial cell output (at time step = 0) [bS, nUn] // Wx input-to-hidden weights, [iS, 3*nUn] // Wh hidden-to-hidden weights, [nUn, 3*nUn] // b biases, [3*nUn] // h is cell outputs at each time step [time, bS, nUn] const int time = x->sizeAt(0); NDArray ht_1(*h0); // loop through time steps for (int t = 0; t < time; ++t) { auto xt = (*x)({t,t+1, 0,0, 0,0}); auto ht = (*h)({t,t+1, 0,0, 0,0}); //helpers::gruCell(&xt, &ht_1, Wx, Wh, b, &ht); //ht_1.assign(ht); } } ////////////////////////////////////////////////////////////////////////// void gruCellBP(nd4j::LaunchContext * context, const NDArray* x, const NDArray* h0, const NDArray* Wx, const NDArray* Wh, const NDArray* b, const NDArray* dLdh, const NDArray* dLdWx0, const NDArray* dLdWh0, const NDArray* dLdb0, NDArray* dLdx, NDArray* dLdh0, NDArray* dLdWx, NDArray* dLdWh, NDArray* dLdb) { // x input [bS, iS] // h0 previous cell output [bS, nUn], that is at previous time step t-1 // Wx input-to-hidden weights, [iS, 3*nUn] // Wh hidden-to-hidden weights, [nUn, 3*nUn] // b biases, [3*nUn] // dLdh gradient wrt output, [bS,nUn], that is epsilon_next // dLdWx0 gradient wrt Wx at previous time step, [iS, 3*nUn] // dLdWh0 gradient wrt Wh at previous time step, [nUn, 3*nUn] // dLdb0 gradient wrt b at previous time step, [3*nUn] // dLdx gradient wrt x, [bS, iS], that is epsilon // dLdh0 gradient wrt h0, [bS, nUn] // dLdWx gradient wrt Wx, [iS, 3*nUn] // dLdWh gradient wrt Wh, [nUn, 3*nUn] // dLdb gradient wrt b at previous time step, [3*nUn] // h is current cell output [bS, nUn], that is at current time step t const int nUn = h0->sizeAt(1); // ***** feed forward step ***** // // gates = sigmoid(x*Wx + h0*Wh + b) auto gates = sigmoid(mmul(*x, (*Wx)({0,0, 0,2*nUn})) + mmul(*h0, (*Wh)({0,0, 0,2*nUn})) + (*b)({0,2*nUn})); // [bS, 2*nUn] + [bS, 2*nUn] + [1, 2*nUn] = [bS, 2*nUn] // reset gate auto r = gates({0,0, 0, nUn}); // [bS, nUn] // update gate auto u = gates({0,0, nUn, 2*nUn}); // [bS, nUn] // ◦ means element-wise product or so called Hadamard product // n = tanh(x*Wx + (r◦h0)*Wh + b) auto n = tanh(mmul(*x, (*Wx)({0,0, 2*nUn,3*nUn})) + mmul((*h0)*r, (*Wh)({0,0, 2*nUn,3*nUn})) + (*b)({2*nUn,3*nUn})); // [bS, nUn] // ***** back prop step ***** // auto Wxr = (*Wx)({0,0, 0, nUn}); auto Wxu = (*Wx)({0,0, nUn, 2*nUn}); auto Wxn = (*Wx)({0,0, 2*nUn,3*nUn}); auto Whr = (*Wh)({0,0, 0, nUn}); auto Whu = (*Wh)({0,0, nUn, 2*nUn}); auto Whn = (*Wh)({0,0, 2*nUn,3*nUn}); auto WxrT = Wxr.transpose(); auto WxuT = Wxu.transpose(); auto WxnT = Wxn.transpose(); auto WhrT = Whr.transpose(); auto WhuT = Whu.transpose(); auto WhnT = Whn.transpose(); auto xT = x->transpose(); auto h0T = h0->transpose(); auto dLdWxr = (*dLdWx)({0,0, 0, nUn}); auto dLdWxu = (*dLdWx)({0,0, nUn, 2*nUn}); auto dLdWxn = (*dLdWx)({0,0, 2*nUn,3*nUn}); auto dLdWhr = (*dLdWh)({0,0, 0, nUn}); auto dLdWhu = (*dLdWh)({0,0, nUn, 2*nUn}); auto dLdWhn = (*dLdWh)({0,0, 2*nUn,3*nUn}); auto dLdbr = (*dLdb)({0, nUn}); auto dLdbu = (*dLdb)({nUn, 2*nUn}); auto dLdbn = (*dLdb)({2*nUn,3*nUn}); auto dhdu = *h0 - n; // [bS, nUn] auto dhdn = 1.f - u; // [bS, nUn] auto dSigdu = u * (1.f - u); // [bS, nUn] auto dSigdr = r * (1.f - r); // [bS, nUn] auto dActdn = 1.f - n * n; // [bS, nUn] auto dndr = mmul(dActdn * (*h0), WhnT); auto drdh0 = mmul(dSigdr, WhrT); auto dLdn = (*dLdh) * dhdn; auto dLdu = (*dLdh) * dhdu; auto dLdr = dLdn * dndr; dLdx->assign( mmul(dLdu * dSigdu, WxuT) + mmul(dLdr * dSigdr, WxrT) + mmul(dLdn * dActdn, WxnT) ); // [bS,iS] dLdh0->assign( mmul(dLdu * dSigdu, WhuT) + mmul(dLdn * dActdn * (r + drdh0), WhnT) + (*dLdh)*u ); // [bS,nUn] dLdWxr.assign( mmul(xT, dSigdr * dLdr) ); // [iS,nUn] dLdWhr.assign( mmul(h0T, dSigdr * dLdr) ); // [nUn,nUn] dLdWxu.assign( mmul(xT, dSigdu * dLdu) ); // [iS,nUn] dLdWhu.assign( mmul(h0T, dSigdu * dLdu) ); // [nUn,nUn] dLdWxn.assign( mmul(xT, dActdn * dLdn) ); // [iS,nUn] dLdWhn.assign( mmul((r*(*h0)).transpose(), dActdn * dLdn) ); // [nUn,nUn] dLdbr.assign( (dSigdr * dLdr).reduceAlongDims(reduce::Sum, {0})); // [nUn] dLdbu.assign( (dSigdu * dLdu).reduceAlongDims(reduce::Sum, {0})); // [nUn] dLdbn.assign( (dActdn * dLdn).reduceAlongDims(reduce::Sum, {0})); // [nUn] if(dLdWx0 != nullptr) *dLdWx += *dLdWx0; if(dLdWh0 != nullptr) *dLdWh += *dLdWh0; if(dLdb0 != nullptr) *dLdb += *dLdb0; } // ////////////////////////////////////////////////////////////////////////// // FIXME - gruTimeLoopBP is not correct // template // void gruTimeLoopBP(const std::vector*>& inArrs, const std::vector*>& outArrs) { // NDArray* x = inArrs[0]; // input [time, bS, iS] // NDArray* hi = inArrs[1]; // previous/initial cell output [bS, nUn], that is at previous time step t-1 // NDArray* Wx = inArrs[2]; // input-to-hidden weights, [iS, 3*nUn] // NDArray* Wh = inArrs[3]; // hidden-to-hidden weights, [nUn, 3*nUn] // NDArray* b = inArrs[4]; // biases, [3*nUn] // NDArray* dLdh = inArrs[5]; // gradient wrt output, [time, bS, nUn], that is epsilon_next // NDArray* dLdx = outArrs[0]; // gradient wrt x, [time, bS, iS], that is epsilon // NDArray* dLdhi = outArrs[1]; // gradient wrt hi, [bS, nUn] // NDArray* dLdWx = outArrs[2]; // gradient wrt Wx, [iS, 3*nUn] // NDArray* dLdWh = outArrs[3]; // gradient wrt Wh, [nUn, 3*nUn] // NDArray* dLdb = outArrs[4]; // gradient wrt b, [3*nUn] // const Nd4jLong time = x->sizeAt(0); // const Nd4jLong bS = x->sizeAt(1); // const Nd4jLong iS = x->sizeAt(2); // const Nd4jLong nUn = hi->sizeAt(1); // NDArray h(hi->ordering(), {time, bS, nUn}); // feed forward output // // first step, time = 0, feed forward // NDArray x0 = (*x)({{0,1}, {}, {}}); // NDArray h0 = h({{0,1}, {}, {}}); // helpers::gruCell({&x0, hi, Wx, Wh, b}, &h0); // // first step, time = 0, back prop // NDArray dLdx0 = (*dLdx)({{0,1}, {}, {}}); // NDArray dLdh0 = (*dLdh)({{0,1}, {}, {}}); // helpers::gruCellBP({&x0, hi, Wx, Wh, b, &dLdh0, nullptr, nullptr, nullptr}, {&dLdx0, dLdhi, dLdWx, dLdWh, dLdb}); // // loop through the rest time steps // for (Nd4jLong t = time-1; t > 0; --t) { // for (Nd4jLong t = 1; t < time; ++t) { // NDArray xt = (*x)({{t,t+1}, {}, {}}); // NDArray ht = h({{t,t+1}, {}, {}}); // NDArray ht_1 = h({{t-1,t}, {}, {}}); // NDArray dLdxt = (*dLdx)({{t,t+1}, {}, {}}); // NDArray dLdht = (*dLdh)({{t,t+1}, {}, {}}); // NDArray dLdWxt_1 = dLdWx; // NDArray dLdWht_1 = dLdWh; // NDArray dLdbt_1 = dLdb; // // feed forward, calculation of ht // helpers::gruCell({&xt, &ht_1, Wx, Wh, b}, &ht); // // back prop // helpers::gruCellBP({&xt, &ht_1, Wx, Wh, b, &dLdht, &dLdWxt_1, &dLdWht_1, &dLdbt_1}, {&dLdxt, nullptr, dLdWx, dLdWh, dLdb}); // } // } } } }