deeplearning4j--deeplearning4j
wehub-resource-sync
1362 行
55 KiB
C++
1362 行
55 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)
|
||
//
|
||
|
||
// implementation of operation for LSTM cell with peep hole connections:
|
||
// http://www.bioinf.jku.at/publications/older/2604.pdf
|
||
// S. Hochreiter and J. Schmidhuber. "Long Short-Term Memory". Neural Computation, 9(8):1735-1780, 1997.
|
||
// and
|
||
// https://research.google.com/pubs/archive/43905.pdf
|
||
// Hasim Sak, Andrew Senior, and Francoise Beaufays. "Long short-term memory recurrent neural network architectures for
|
||
// large scale acoustic modeling." INTERSPEECH, 2014.
|
||
|
||
|
||
#include <system/op_boilerplate.h>
|
||
|
||
|
||
#if NOT_EXCLUDED(OP_lstmLayer)
|
||
|
||
#include <execution/Threads.h>
|
||
#include <helpers/MmulHelper.h>
|
||
#include <helpers/ShapeUtils.h>
|
||
#include <ops/declarable/helpers/activations.h>
|
||
#include <ops/declarable/helpers/lstmLayer.h>
|
||
|
||
|
||
namespace sd {
|
||
namespace ops {
|
||
namespace helpers {
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
static void applyActivation(NDArray *x, const int opId, const float alpha, const float beta, NDArray* z) {
|
||
switch (opId) {
|
||
case 0:
|
||
x->applyTransform(transform::Tanh, z);
|
||
break;
|
||
case 1:
|
||
x->applyScalar(scalar::RELU, 0, z);
|
||
break;
|
||
case 2:
|
||
x->applyTransform(transform::Sigmoid, z);
|
||
break;
|
||
case 3: {
|
||
ExtraArguments args({static_cast<double>(alpha), static_cast<double>(beta)});
|
||
x->applyTransform(transform::Affine, z, &args);
|
||
break;
|
||
}
|
||
case 4:
|
||
x->applyScalar(scalar::LeakyRELU, alpha, z);
|
||
break;
|
||
case 5:
|
||
thresholdRelu(x->getContext(), x, alpha, z);
|
||
break;
|
||
case 6: {
|
||
ExtraArguments args({static_cast<double>(alpha), static_cast<double>(beta)});
|
||
x->applyTransform(transform::ScaledTanh, z, &args);
|
||
break;
|
||
}
|
||
case 7:
|
||
x->applyTransform(transform::HardSigmoid, z);
|
||
break;
|
||
case 8:
|
||
x->applyScalar(scalar::ELU, alpha, z);
|
||
break;
|
||
case 9:
|
||
x->applyTransform(transform::SoftSign, z);
|
||
break;
|
||
case 10:
|
||
x->applyTransform(transform::SoftPlus, z);
|
||
break;
|
||
default:
|
||
THROW_EXCEPTION("LSTM_LAYER operation: wrong id number of activation !");
|
||
}
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
static void activationDeriv(NDArray* x, const int opId, const float alpha, const float beta, NDArray* z) {
|
||
switch (opId) {
|
||
case 0:
|
||
x->applyTransform(transform::TanhDerivative, z);
|
||
break;
|
||
case 1:
|
||
x->applyScalar(scalar::RELUDerivative, 0, z);
|
||
break;
|
||
case 2:
|
||
x->applyTransform(transform::SigmoidDerivative, z);
|
||
break;
|
||
case 3: {
|
||
*z = alpha;
|
||
break;
|
||
}
|
||
case 4:
|
||
x->applyScalar(scalar::LeakyRELUDerivative, alpha, z);
|
||
break;
|
||
case 5:
|
||
thresholdReluDerivative(x->getContext(), x, alpha, z, z);
|
||
break;
|
||
case 6: {
|
||
auto func = PRAGMA_THREADS_FOR {
|
||
for (sd::LongType i = start; i < stop; ++i) {
|
||
auto val = beta * x->e<float>(i);
|
||
z->p<float>(
|
||
i, alpha * beta * (1.f - sd::math::sd_tanh<float, float>(val) * sd::math::sd_tanh<float, float>(val)));
|
||
}
|
||
};
|
||
samediff::Threads::parallel_for(func, 0, x->lengthOf());
|
||
break;
|
||
}
|
||
case 7:
|
||
x->applyTransform(transform::HardSigmoidDerivative, z);
|
||
break;
|
||
case 8:
|
||
x->applyScalar(scalar::ELUDerivative, alpha, z);
|
||
break;
|
||
case 9:
|
||
x->applyTransform(transform::SoftSignDerivative, z);
|
||
break;
|
||
case 10: {
|
||
auto func = PRAGMA_THREADS_FOR {
|
||
for (sd::LongType i = start; i < stop; ++i) {
|
||
auto val = sd::math::sd_exp<float, float>(x->e<float>(i));
|
||
z->p<float>(i, val / (1.f + val));
|
||
}
|
||
};
|
||
samediff::Threads::parallel_for(func, 0, x->lengthOf());
|
||
break;
|
||
}
|
||
default:
|
||
THROW_EXCEPTION("LSTM_LAYER operation: wrong id number of activation !");
|
||
}
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
// FIXME - derivative undefined when not-clipped c has element/elements equal to -clipVal or clipVal
|
||
static void clipDeriv(const float clipVal, NDArray& c, NDArray& z0, NDArray& z1, NDArray& z2, NDArray& z3) {
|
||
if (clipVal == 0) return;
|
||
|
||
auto func = PRAGMA_THREADS_FOR {
|
||
for (sd::LongType i = start; i < stop; ++i) {
|
||
const auto val = c.e<float>(i);
|
||
if (val == -clipVal || val == clipVal) {
|
||
z0.p<float>(i, 0.f);
|
||
z1.p<float>(i, 0.f);
|
||
z2.p<float>(i, 0.f);
|
||
z3.p<float>(i, 0.f);
|
||
}
|
||
}
|
||
};
|
||
samediff::Threads::parallel_for(func, 0, c.lengthOf());
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
static NDArray *tensorAlongTimeBatchDims(NDArray& arr, const int dataFormat, const int t1, const int t2,
|
||
const int b1, const int b2) {
|
||
if (dataFormat == 0 || dataFormat == 3) return arr({t1, t2, b1, b2, 0, 0}); // TNS: [sL, bS, nIn]
|
||
|
||
if (dataFormat == 1) return arr({b1, b2, t1, t2, 0, 0}); // NTS: [bS, sL ,nIn]
|
||
|
||
return arr({b1, b2, 0, 0, t1, t2}); // NST: [bS, nIn, sL]
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
static SD_INLINE int getBatchTimeTotalIndex(const int dataFormat, const int sL, const int bS, const int t,
|
||
const int b) {
|
||
if (dataFormat == 0 || dataFormat == 3) return t * bS + b; // TNS: shape [sL, bS, nIn]
|
||
|
||
return b * sL + t; // NTS, NST: shape [bS, sL, nIn], [bS, nIn, sL]
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void lstmLayerCell(NDArray* x, NDArray* Wx, NDArray* Wr, NDArray* b, NDArray* hI,
|
||
NDArray* cI, NDArray* Wp, const std::vector<float>& params, NDArray* h, NDArray* c) {
|
||
// * -> means element-wise multiplication
|
||
// × -> means matrix multiplication
|
||
|
||
/************************ THIS IS NOT OPTIMAZED CODE ***********************************/
|
||
/** the objective is to provide math-readable code **/
|
||
|
||
// equations (no peephole connections)
|
||
// it = σ(Wxi × xt + Wri × ht-1 + bi)
|
||
// ft = σ(Wxf × xt + Wrf × ht-1 + bf)
|
||
// c't = tanh(Wxc × xt + Wrc × ht-1 + bc)
|
||
// ct = ft * ct-1 + it * c't
|
||
// ot = σ(Wxo × xt + Wro × ht-1 + bo)
|
||
// ht = ot * tanh(ct)
|
||
|
||
// equations (peephole connections are present)
|
||
// it = σ(Wxi × xt + Wri × ht-1 + Wpi * ct-1 + bi)
|
||
// ft = σ(Wxf × xt + Wrf × ht-1 + Wpf * ct-1 + bf)
|
||
// c't = tanh(Wxc × xt + Wrc × ht-1 + bc)
|
||
// ct = ft * ct-1 + it * c't
|
||
// ot = σ(Wxo × xt + Wro × ht-1 + Wpo * ct + bo)
|
||
// ht = ot * tanh(ct)
|
||
|
||
// IDs for activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded relu, 6=scaled tanh, 7=hard
|
||
// sigmoid, 8=ELU, 9=softsign, 10=softplus
|
||
|
||
// params[0] - dataFormat, ignore
|
||
// params[1] - directionMode, ignore
|
||
// params[2] - cell clipping value, if it = 0 then do not apply clipping
|
||
|
||
// params[3] - activation ID for input (i), forget (f) and output (o) gates
|
||
// params[4] - alpha value for gates activation
|
||
// params[5] - beta value for gates activation
|
||
|
||
// params[6] - activation ID for cell state (c)
|
||
// params[7] - alpha value for cell state activation
|
||
// params[8] - beta value for cell state activation
|
||
|
||
// params[9] - activation ID for output (h)
|
||
// params[10] - alpha value for output activation
|
||
// params[11] - beta value for output activation
|
||
|
||
// INPUTS:
|
||
// x - current input at time t, [bS, nIn] or [nIn] if seqLen != nullptr
|
||
// Wx - input weights [nIn, 4*nOut]
|
||
// Wr - recurrent weights [nOut, 4*nOut]
|
||
// b - biases [4*nOut], optional, may be nullptr
|
||
// hI - (ht-1) previous (initial) output at time t-1, optional may be nullptr, [bS, nOut] or [nOut] if seqLen !=
|
||
// nullptr cI - (ct-1) previous (initial) cell state at time t-1, optional may be nullptr, [bS, nOut] or [nOut] if
|
||
// seqLen != nullptr Wp - peephole weights [3*nOut], optional, may be nullptr
|
||
|
||
// OUTPUTS:
|
||
// h - current output, that is at current time step t, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// c - current cell state, that is at current time step t, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
|
||
// !!! dimension 4*nOut implies order it, ft, c't, ot
|
||
// !!! dimension 3*nOut implies order it, ft, ot
|
||
|
||
const sd::LongType nOut = Wx->sizeAt(-1) / 4;
|
||
|
||
NDArray *xMulWx = mmul(*x, *Wx); // [bs, nIn] * [nIn, 4*nOut] = [bS, 4*nOut] or [nIn] * [nIn, 4*nOut] = [4*nOut]
|
||
NDArray *hIMulWr = mmul(*hI, *Wr); // [bs, nOut] * [nOut, 4*nOut] = [bS, 4*nOut] or [nOut] * [nOut, 4*nOut] = [4*nOut]
|
||
auto* z = (*xMulWx) + (*hIMulWr);
|
||
delete xMulWx;
|
||
delete hIMulWr;
|
||
|
||
// add biases if they are given
|
||
if (b != nullptr) *z += *b; // broadcast [bS, 4*nOut](or[4*nOut]) + [4*nOut] = [bS, 4*nOut]
|
||
|
||
auto zi = x->rankOf() == 1 ? (*z)({0, nOut}) : (*z)({0, 0, 0, nOut}); // input gate it, [bS, nOut](or[nOut])
|
||
auto zf = x->rankOf() == 1 ? (*z)({nOut, 2 * nOut}) : (*z)({0, 0, nOut, 2 * nOut}); // forget gate ft, [bS, nOut](or[nOut])
|
||
auto zg = x->rankOf() == 1 ? (*z)({2 * nOut, 3 * nOut})
|
||
: (*z)({0, 0, 2 * nOut, 3 * nOut}); // cell gate c't, [bS, nOut](or[nOut])
|
||
auto zo = x->rankOf() == 1 ? (*z)({3 * nOut, 4 * nOut})
|
||
: (*z)({0, 0, 3 * nOut, 4 * nOut}); // output gate ot, [bS, nOut](or[nOut])
|
||
|
||
// peephole connections for input and forget gates
|
||
if (Wp != nullptr) {
|
||
NDArray *wpFirst = (*Wp)({0, nOut});
|
||
NDArray *wpSecond = (*Wp)({nOut, 2 * nOut});
|
||
NDArray *cIMulWpFirst = (*cI) * (*wpFirst); // broadcast: [bS, nOut] * [nOut] = [bS, nOut](or[nOut])
|
||
NDArray *cIMulWpSecond = (*cI) * (*wpSecond); // broadcast: [bS, nOut] * [nOut] = [bS, nOut](or[nOut])
|
||
*zi += *cIMulWpFirst;
|
||
*zf += *cIMulWpSecond;
|
||
delete cIMulWpFirst;
|
||
delete cIMulWpSecond;
|
||
delete wpFirst;
|
||
delete wpSecond;
|
||
}
|
||
|
||
applyActivation(zi, params[3], params[4], params[5], zi); // inplace
|
||
applyActivation(zf, params[3], params[4], params[5], zf); // inplace
|
||
applyActivation(zg, params[6], params[7], params[8], zg); // inplace
|
||
|
||
NDArray *zfMulCI = (*zf) * (*cI); // [bS, nOut] * [bS, nOut] = [bS, nOut](or[nOut])
|
||
NDArray *ziMulZg = (*zi) * (*zg); // [bS, nOut] * [bS, nOut] = [bS, nOut](or[nOut])
|
||
NDArray *cAssign = (*zfMulCI) + (*ziMulZg);
|
||
c->assign(cAssign);
|
||
delete zfMulCI;
|
||
delete ziMulZg;
|
||
delete cAssign;
|
||
|
||
// if clipping value is non-zero then cell state is clipped by this value prior to the cell output activation
|
||
if (params[2] != 0) c->applyScalar(scalar::LstmClip, params[2], c);
|
||
|
||
// peephole connections for output gate
|
||
if (Wp != nullptr) {
|
||
NDArray *wpThird = (*Wp)({2 * nOut, 3 * nOut});
|
||
NDArray *cMulWpThird = (*c) * (*wpThird); // broadcast: [bS, nOut] * [nOut] = [bS, nOut](or[nOut])
|
||
*zo += *cMulWpThird;
|
||
delete cMulWpThird;
|
||
delete wpThird;
|
||
}
|
||
|
||
applyActivation(zo, params[3], params[4], params[5], zo);
|
||
|
||
applyActivation(c, params[9], params[10], params[11], h);
|
||
*h *= *zo; // [bS, nOut] * [bS, nOut](or[nOut])
|
||
|
||
delete z;
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
// this auxiliary ff should be running before backprop
|
||
void lstmLayerCell(NDArray* x, NDArray* Wx, NDArray* Wr, NDArray* b, NDArray* hI,
|
||
NDArray* cI, NDArray* Wp, const std::vector<float>& params, NDArray* z, NDArray* a,
|
||
NDArray* h, NDArray* c) {
|
||
// z - zi, zf, zg, zo
|
||
// a - i, f, g, o
|
||
|
||
const sd::LongType nOut = Wx->sizeAt(-1) / 4;
|
||
|
||
NDArray *xMulWx = mmul(*x, *Wx); // [bs, nIn] * [nIn, 4*nOut] = [bS, 4*nOut] or [nIn] * [nIn, 4*nOut] = [4*nOut]
|
||
NDArray *hIMulWr = mmul(*hI, *Wr); // [bs, nOut] * [nOut, 4*nOut] = [bS, 4*nOut] or [nOut] * [nOut, 4*nOut] = [4*nOut]
|
||
NDArray *zAssign = (*xMulWx) + (*hIMulWr);
|
||
z->assign(zAssign);
|
||
delete xMulWx;
|
||
delete hIMulWr;
|
||
delete zAssign;
|
||
|
||
// add biases if they are given
|
||
if (b != nullptr) *z += *b; // broadcast [bS, 4*nOut](or[4*nOut]) + [4*nOut] = [bS, 4*nOut]
|
||
|
||
auto zi = x->rankOf() == 1 ? (*z)({0, nOut}) : (*z)({0, 0, 0, nOut}); // input gate it, [bS, nOut](or[nOut])
|
||
auto zf =
|
||
x->rankOf() == 1 ? (*z)({nOut, 2 * nOut}) : (*z)({0, 0, nOut, 2 * nOut}); // forget gate ft, [bS, nOut](or[nOut])
|
||
auto zg = x->rankOf() == 1 ? (*z)({2 * nOut, 3 * nOut})
|
||
: (*z)({0, 0, 2 * nOut, 3 * nOut}); // cell gate c't, [bS, nOut](or[nOut])
|
||
auto zo = x->rankOf() == 1 ? (*z)({3 * nOut, 4 * nOut})
|
||
: (*z)({0, 0, 3 * nOut, 4 * nOut}); // output gate ot, [bS, nOut](or[nOut])
|
||
|
||
auto i = x->rankOf() == 1 ? (*a)({0, nOut}) : (*a)({0, 0, 0, nOut}); // input gate it, [bS, nOut](or[nOut])
|
||
auto f =
|
||
x->rankOf() == 1 ? (*a)({nOut, 2 * nOut}) : (*a)({0, 0, nOut, 2 * nOut}); // forget gate ft, [bS, nOut](or[nOut])
|
||
auto g = x->rankOf() == 1 ? (*a)({2 * nOut, 3 * nOut})
|
||
: (*a)({0, 0, 2 * nOut, 3 * nOut}); // cell gate c't, [bS, nOut](or[nOut])
|
||
auto o = x->rankOf() == 1 ? (*a)({3 * nOut, 4 * nOut})
|
||
: (*a)({0, 0, 3 * nOut, 4 * nOut}); // output gate ot, [bS, nOut](or[nOut])
|
||
|
||
// peephole connections for input and forget gates
|
||
if (Wp != nullptr) {
|
||
NDArray *wpFirst = (*Wp)({0, nOut});
|
||
NDArray *wpSecond = (*Wp)({nOut, 2 * nOut});
|
||
NDArray *cIMulWpFirst = (*cI) * (*wpFirst); // broadcast: [bS, nOut] * [nOut] = [bS, nOut](or[nOut])
|
||
NDArray *cIMulWpSecond = (*cI) * (*wpSecond); // broadcast: [bS, nOut] * [nOut] = [bS, nOut](or[nOut])
|
||
*zi += *cIMulWpFirst;
|
||
*zf += *cIMulWpSecond;
|
||
delete cIMulWpFirst;
|
||
delete cIMulWpSecond;
|
||
delete wpFirst;
|
||
delete wpSecond;
|
||
}
|
||
|
||
applyActivation(zi, params[3], params[4], params[5], i);
|
||
applyActivation(zf, params[3], params[4], params[5], f);
|
||
applyActivation(zg, params[6], params[7], params[8], g);
|
||
|
||
NDArray *fMulCI = (*f) * (*cI); // [bS, nOut] * [bS, nOut] = [bS, nOut](or[nOut])
|
||
NDArray *iMulG = (*i) * (*g); // [bS, nOut] * [bS, nOut] = [bS, nOut](or[nOut])
|
||
NDArray *cAssign = (*fMulCI) + (*iMulG);
|
||
c->assign(cAssign);
|
||
delete fMulCI;
|
||
delete iMulG;
|
||
delete cAssign;
|
||
|
||
// if clipping value is non-zero then cell state is clipped by this value prior to the cell output activation
|
||
if (params[2] != 0) c->applyScalar(scalar::LstmClip, params[2], c);
|
||
|
||
// peephole connections for output gate
|
||
if (Wp != nullptr) {
|
||
NDArray *wpThird = (*Wp)({2 * nOut, 3 * nOut});
|
||
NDArray *cMulWpThird = (*c) * (*wpThird); // broadcast: [bS, nOut] * [nOut] = [bS, nOut](or[nOut])
|
||
*zo += *cMulWpThird;
|
||
delete cMulWpThird;
|
||
delete wpThird;
|
||
}
|
||
|
||
applyActivation(zo, params[3], params[4], params[5], o);
|
||
|
||
applyActivation(c, params[9], params[10], params[11], h);
|
||
*h *= *o; // [bS, nOut] * [bS, nOut](or[nOut])
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void lstmLayerCellBp(NDArray* x, NDArray* Wx, NDArray* Wr, NDArray* b, NDArray* hI,
|
||
NDArray* cI, NDArray* Wp, NDArray* dLdh, NDArray* dLdhL,
|
||
NDArray* dLdcL, NDArray* z, NDArray* a, NDArray* c,
|
||
const std::vector<float>& params, NDArray* dLdx, NDArray* dLdWx, NDArray* dLdWr, NDArray* dLdhI,
|
||
NDArray* dLdcI, NDArray* dLdb, NDArray* dLdWp) {
|
||
/************************ THIS IS NOT OPTIMAZED CODE ***********************************/
|
||
/** the objective is to provide math-readable code **/
|
||
|
||
// equations (no peephole connections)
|
||
// zi = x × Wxi + hI × Wri + bi
|
||
// zf = x × Wxf + hI × Wrf + bf
|
||
// zg = x × Wxg + hI × Wrg + bg
|
||
// zo = x × Wxo + hI × Wro + bo
|
||
// i = act(zi)
|
||
// f = act(zf)
|
||
// g = actC(zg)
|
||
// o = act(zo)
|
||
// c = clip(f * cI + i * g)
|
||
// h = o * actH(c)
|
||
|
||
// equations (peephole connections are present)
|
||
// zi = x × Wxi + hI × Wri + cI * Wpi + bi
|
||
// zf = x × Wxf + hI × Wrf + cI * Wpf + bf
|
||
// zg = x × Wxg + hI × Wrg + bg
|
||
// zo = x × Wxo + hI × Wro + c * Wpo + bo
|
||
// i = act(zi)
|
||
// f = act(zf)
|
||
// g = actC(zg)
|
||
// o = act(zo)
|
||
// c = clip(f * cI + i * g)
|
||
// h = o * actH(c)
|
||
|
||
// IDs for activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded relu, 6=scaled tanh, 7=hard
|
||
// sigmoid, 8=ELU, 9=softsign, 10=softplus
|
||
|
||
// params[0] - dataFormat, ignore
|
||
// params[1] - directionMode, ignore
|
||
// params[2] - cell clipping value, if it = 0 then do not apply clipping
|
||
|
||
// params[3] - activation ID for input (i), forget (f) and output (o) gates
|
||
// params[4] - alpha value for gates activation
|
||
// params[5] - beta value for gates activation
|
||
|
||
// params[6] - activation ID for cell state (c)
|
||
// params[7] - alpha value for cell state activation
|
||
// params[8] - beta value for cell state activation
|
||
|
||
// params[9] - activation ID for output (h)
|
||
// params[10] - alpha value for output activation
|
||
// params[11] - beta value for output activation
|
||
|
||
// INPUTS:
|
||
// x - current input at time t, [bS, nIn] or [nIn] if seqLen != nullptr
|
||
// Wx - input weights [nIn, 4*nOut]
|
||
// Wr - recurrent weights [nOut, 4*nOut]
|
||
// b - biases [4*nOut], optional, may be nullptr
|
||
// hI - (ht-1) previous (initial) output at time t-1, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// cI - (ct-1) previous (initial) cell state at time t-1, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// Wp - peephole weights [3*nOut], optional, may be nullptr
|
||
// dLdh - loss derivative with respect to h at each time step, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// dLdhL - loss derivative with respect to h at last time step, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// dLdcL - loss derivative with respect to c at last time step, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// z - zi,zf,zg,zo taken from ff outputs to reduce amount of calculations in bp, [bS, 4*nOut]
|
||
// a - i,f,g,o taken from ff outputs to reduce amount of calculations in bp, [bS, 4*nOut]
|
||
// c - taken from ff outputs to reduce amount of calculations in bp, [bS, nOut]
|
||
|
||
// OUTPUTS:
|
||
// dLdx - loss derivative with respect to x, [bS, nIn] or [nIn] if seqLen != nullptr
|
||
// dLdWx - loss derivative with respect to Wx, [nIn, 4*nOut]
|
||
// dLdWr - loss derivative with respect to Wr, [nOut, 4*nOut]
|
||
// dLdb - loss derivative with respect to b, optional, may be nullptr, [4*nOut]
|
||
// dLdhI - loss derivative with respect to hI, optional may be nullptr, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// dLdcI - loss derivative with respect to cI, optional may be nullptr, [bS, nOut] or [nOut] if seqLen != nullptr
|
||
// dLdWp - loss derivative with respect to Wp, optional, may be nullptr, [3*nOut]
|
||
|
||
// !!! dimension 4*nOut implies order i, f, g, o
|
||
// !!! dimension 3*nOut implies order i, f, o
|
||
|
||
// dhdc = o*tanhDeriv + Wp ? tanh(c)*dodzo*dzodc : 0 [bS, nOut]
|
||
// dcdcI = f + Wp ? dcdzi*dzidcI + dcdzf*dzfdcI : 0 [bS, nOut]
|
||
|
||
// dLdhI += dLdh; [bS, nOut]
|
||
// dLdcI += dLdhI * dhdc; [bS, nOut]
|
||
|
||
// dLdzi = dLdcI*dcdi*didzi; [bS, nOut](or[nOut])
|
||
// dLdzf = dLdcI*dcdf*dfdzf; [bS, nOut](or[nOut])
|
||
// dLdzg = dLdcI*dcdg*dgdzg; [bS, nOut](or[nOut])
|
||
// dLdzo = dLdhI*dhdo*dodzo; [bS, nOut](or[nOut])
|
||
|
||
// dLdx = dLdzi×WxiT + dLdzf×WxfT + dLdzg×WxgT + dLdzo×WxoT, [bS, nIn]
|
||
// dLdhI = dLdzi×WriT + dLdzf×WrfT + dLdzg×WrgT + dLdzo×WroT, [bS, nOut]
|
||
// dLdcI = dLdcI*dcdcI, [bS, nOut]
|
||
|
||
// dLdWxi = xT×dLdzi [nIn, bS] x [bS, nOut] = [nIn, nOut]
|
||
// dLdWxf = xT×dLdzf [nIn, bS] x [bS, nOut] = [nIn, nOut]
|
||
// dLdWxg = xT×dLdzg [nIn, bS] x [bS, nOut] = [nIn, nOut]
|
||
// dLdWxo = xT×dLdzo [nIn, bS] x [bS, nOut] = [nIn, nOut]
|
||
|
||
// dLdWri = hIT×dLdzi [nOut, bS] x [bS, nOut] = [nOut, nOut]
|
||
// dLdWrf = hIT×dLdzf [nOut, bS] x [bS, nOut] = [nOut, nOut]
|
||
// dLdWrg = hIT×dLdzg [nOut, bS] x [bS, nOut] = [nOut, nOut]
|
||
// dLdWro = hIT×dLdzo [nOut, bS] x [bS, nOut] = [nOut, nOut]
|
||
|
||
// dLdbi = dLdzi.reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
// dLdbf = dLdzf.reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
// dLdbg = dLdzg.reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
// dLdbo = dLdzo.reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
|
||
// dLdWpi = (dLdzi*cI).reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
// dLdWpf = (dLdzf*cI).reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
// dLdWpo = (dLdzo*c) .reduce_sum_along_0_axis [bS, nOut] -> reduce -> [nOut]
|
||
|
||
const sd::LongType nOut = Wx->sizeAt(-1) / 4;
|
||
const sd::LongType nIn = x->sizeAt(-1);
|
||
|
||
NDArray *zi = x->rankOf() == 1 ? (*z)({0, nOut}) : (*z)({0, 0, 0, nOut}); // input gate i, [bS, nOut](or[nOut])
|
||
NDArray *zf =
|
||
x->rankOf() == 1 ? (*z)({nOut, 2 * nOut}) : (*z)({0, 0, nOut, 2 * nOut}); // forget gate f, [bS, nOut](or[nOut])
|
||
NDArray *zg = x->rankOf() == 1 ? (*z)({2 * nOut, 3 * nOut})
|
||
: (*z)({0, 0, 2 * nOut, 3 * nOut}); // cell gate g, [bS, nOut](or[nOut])
|
||
NDArray *zo = x->rankOf() == 1 ? (*z)({3 * nOut, 4 * nOut})
|
||
: (*z)({0, 0, 3 * nOut, 4 * nOut}); // output gate o, [bS, nOut](or[nOut])
|
||
|
||
NDArray *i = x->rankOf() == 1 ? (*a)({0, nOut}) : (*a)({0, 0, 0, nOut}); // input gate i, [bS, nOut](or[nOut])
|
||
NDArray *f =
|
||
x->rankOf() == 1 ? (*a)({nOut, 2 * nOut}) : (*a)({0, 0, nOut, 2 * nOut}); // forget gate f, [bS, nOut](or[nOut])
|
||
NDArray *g = x->rankOf() == 1 ? (*a)({2 * nOut, 3 * nOut})
|
||
: (*a)({0, 0, 2 * nOut, 3 * nOut}); // cell gate g, [bS, nOut](or[nOut])
|
||
NDArray *o = x->rankOf() == 1 ? (*a)({3 * nOut, 4 * nOut})
|
||
: (*a)({0, 0, 3 * nOut, 4 * nOut}); // output gate o, [bS, nOut](or[nOut])
|
||
|
||
NDArray *zUlike = z->ulike(); // [bS, 4*nOut](or[4*nOut])
|
||
NDArray dLdz = *zUlike; // [bS, 4*nOut](or[4*nOut])
|
||
NDArray *dLdzi = x->rankOf() == 1 ? dLdz({0, nOut}) : dLdz({0, 0, 0, nOut});
|
||
NDArray *dLdzf = x->rankOf() == 1 ? dLdz({nOut, 2 * nOut}) : dLdz({0, 0, nOut, 2 * nOut});
|
||
NDArray *dLdzg = x->rankOf() == 1 ? dLdz({2 * nOut, 3 * nOut}) : dLdz({0, 0, 2 * nOut, 3 * nOut});
|
||
NDArray *dLdzo = x->rankOf() == 1 ? dLdz({3 * nOut, 4 * nOut}) : dLdz({0, 0, 3 * nOut, 4 * nOut});
|
||
|
||
// dcdzi = dcdi*didzi, [bS, nOut](or[nOut])
|
||
activationDeriv(zi, params[3], params[4], params[5], dLdzi); // didzi, inplace
|
||
NDArray *dLdziMulG = (*dLdzi) * (*g); // dcdi = g*clipDeriv
|
||
dLdzi->assign(dLdziMulG);
|
||
delete dLdziMulG;
|
||
|
||
// dcdzf = dcdf*dfdzf, [bS, nOut](or[nOut])
|
||
activationDeriv(zf, params[3], params[4], params[5], dLdzf); // dfdzf, inplace
|
||
NDArray *dLdzfMulCI = (*dLdzf) * (*cI); // dcdf = cI*clipDeriv
|
||
dLdzf->assign(dLdzfMulCI);
|
||
delete dLdzfMulCI;
|
||
|
||
// dcdzg = dcde*dedzg, [bS, nOut](or[nOut])
|
||
activationDeriv(zg, params[6], params[7], params[8], dLdzg); // dgdzg, inplace
|
||
NDArray *dLdzgMulI = (*dLdzg) * (*i); // dcdf = i*clipDeriv
|
||
dLdzg->assign(dLdzgMulI);
|
||
delete dLdzgMulI;
|
||
|
||
// dhdzo = dhdo*dodzo = actH(c)*dodzo, [bS, nOut](or[nOut])
|
||
activationDeriv(zo, params[3], params[4], params[5], dLdzo);
|
||
NDArray *dLdzoUlike = dLdzo->ulike();
|
||
NDArray temp = *dLdzoUlike;
|
||
applyActivation(c, params[9], params[10], params[11], &temp); // actH(c), inplace
|
||
NDArray *dLdzoMulTemp = (*dLdzo) * temp;
|
||
dLdzo->assign(dLdzoMulTemp);
|
||
delete dLdzoMulTemp;
|
||
|
||
// dcdcI
|
||
NDArray *dcdcI = f->dup();
|
||
|
||
// take into account possible deposit from clipping derivative
|
||
clipDeriv(params[2], *c, *dLdzi, *dLdzf, *dLdzg, *dcdcI);
|
||
|
||
// dhdc
|
||
NDArray *cUlike = c->ulike();
|
||
NDArray dhdc = *cUlike;
|
||
activationDeriv(c, params[9], params[10], params[11], &dhdc); // [bS, nOut]
|
||
NDArray *dhdcMulO = dhdc * (*o);
|
||
dhdc.assign(dhdcMulO);
|
||
delete dhdcMulO;
|
||
|
||
if (Wp) {
|
||
NDArray *wpThird = (*Wp)({2 * nOut, 3 * nOut});
|
||
NDArray *wpFirst = (*Wp)({0, nOut});
|
||
NDArray *wpSecond = (*Wp)({nOut, 2 * nOut});
|
||
|
||
NDArray *dLdzoMulWpThird = (*dLdzo) * (*wpThird);
|
||
NDArray *dhdcPlusDLdzoMulWpThird = dhdc + (*dLdzoMulWpThird);
|
||
dhdc.assign(dhdcPlusDLdzoMulWpThird);
|
||
delete dLdzoMulWpThird;
|
||
delete dhdcPlusDLdzoMulWpThird;
|
||
|
||
NDArray *dLdziMulWpFirst = (*dLdzi) * (*wpFirst); // broadcast [bS, nOut] * nOut
|
||
NDArray *dLdzfMulWpSecond = (*dLdzf) * (*wpSecond); // broadcast [bS, nOut] * nOut
|
||
NDArray *sumTemp = (*dLdziMulWpFirst) + (*dLdzfMulWpSecond);
|
||
NDArray *dcdcIPlusSumTemp = (*dcdcI) + (*sumTemp);
|
||
dcdcI->assign(dcdcIPlusSumTemp);
|
||
delete dLdziMulWpFirst;
|
||
delete dLdzfMulWpSecond;
|
||
delete sumTemp;
|
||
delete dcdcIPlusSumTemp;
|
||
|
||
delete wpThird;
|
||
delete wpFirst;
|
||
delete wpSecond;
|
||
}
|
||
|
||
if (dLdh) *dLdhI += *dLdh;
|
||
if (dLdhL) *dLdhI += *dLdhL;
|
||
if (dLdcL) *dLdcI += *dLdcL;
|
||
|
||
NDArray *dLdhIMulDhdc = (*dLdhI) * dhdc;
|
||
*dLdcI += *dLdhIMulDhdc;
|
||
delete dLdhIMulDhdc;
|
||
|
||
NDArray *dLdziMulDLdcI = (*dLdzi) * (*dLdcI); // [bS, nOut](or[nOut])
|
||
dLdzi->assign(dLdziMulDLdcI);
|
||
delete dLdziMulDLdcI;
|
||
|
||
NDArray *dLdzfMulDLdcI = (*dLdzf) * (*dLdcI); // [bS, nOut](or[nOut])
|
||
dLdzf->assign(dLdzfMulDLdcI);
|
||
delete dLdzfMulDLdcI;
|
||
|
||
NDArray *dLdzgMulDLdcI = (*dLdzg) * (*dLdcI); // [bS, nOut](or[nOut])
|
||
dLdzg->assign(dLdzgMulDLdcI);
|
||
delete dLdzgMulDLdcI;
|
||
|
||
NDArray *dLdzoMulDLdhI = (*dLdzo) * (*dLdhI); // [bS, nOut](or[nOut])
|
||
dLdzo->assign(dLdzoMulDLdhI);
|
||
delete dLdzoMulDLdhI;
|
||
|
||
// dLdx
|
||
NDArray *WxT = Wx->transpose();
|
||
MmulHelper::mmul(&dLdz, WxT,
|
||
dLdx); // [bS, 4*nOut] x [4*nOut, nIn] (or [4*nOut] x [4*nOut, nIn]) = [bS, nIn] ( or[nIn] )
|
||
|
||
// dLdhI
|
||
NDArray *WrT = Wr->transpose();
|
||
MmulHelper::mmul(&dLdz, WrT,
|
||
dLdhI); // [bS, 4*nOut] x [4*nOut, nOut] (or [4*nOut] x [4*nOut, nOut]) = [bS, nOut] ( or[nOut] )
|
||
|
||
// dLdcI
|
||
NDArray *dLdcIMulDcdcI = (*dLdcI) * (*dcdcI);
|
||
dLdcI->assign(dLdcIMulDcdcI); // [bS, nOut](or[nOut])
|
||
delete dLdcIMulDcdcI;
|
||
|
||
delete WxT;
|
||
delete WrT;
|
||
|
||
if (x->rankOf() == 1) {
|
||
std::vector<sd::LongType> xShape = {nIn, 1};
|
||
std::vector<sd::LongType> hIShape = {nOut, 1};
|
||
std::vector<sd::LongType> dLdzShape = {1, 4 * nOut};
|
||
NDArray *xT = x->reshape(x->ordering(), xShape); // [nIn] -> [nIn, 1]
|
||
NDArray *hIT = hI->reshape(hI->ordering(), hIShape); // [nOut] -> [nOut, 1]
|
||
NDArray *dLdzR = dLdz.reshape(dLdz.ordering(), dLdzShape); // [nOut] -> [1, 4*nOut]
|
||
|
||
// dLdWx
|
||
NDArray *xTMulDLdzR = mmul(*xT, *dLdzR); // [nIn, 1] x [1, 4*nOut] = [nIn, 4*nOut]
|
||
*dLdWx += *xTMulDLdzR;
|
||
delete xTMulDLdzR;
|
||
|
||
// dLdWr
|
||
NDArray *hITMulDLdzR = mmul(*hIT, *dLdzR); // [nOut, 1] x [1, 4*nOut] = [nOut, 4*nOut]
|
||
*dLdWr += *hITMulDLdzR;
|
||
delete hITMulDLdzR;
|
||
|
||
delete xT;
|
||
delete hIT;
|
||
delete dLdzR;
|
||
} else {
|
||
NDArray *xT = x->transpose();
|
||
NDArray *hIT = hI->transpose();
|
||
|
||
// dLdWx
|
||
NDArray *xTMulDLdz = mmul(*xT, dLdz); // [nIn, bS] x [bS, 4*nOut] = [nIn, 4*nOut]
|
||
*dLdWx += *xTMulDLdz;
|
||
delete xTMulDLdz;
|
||
|
||
// dLdWr
|
||
NDArray *hITMulDLdz = mmul(*hIT, dLdz); // [nOut, bS] x [bS, 4*nOut] = [nOut, 4*nOut]
|
||
*dLdWr += *hITMulDLdz;
|
||
delete hITMulDLdz;
|
||
|
||
delete xT;
|
||
delete hIT;
|
||
}
|
||
|
||
// dLdb
|
||
if (b && x->rankOf() == 1)
|
||
*dLdb += dLdz; // [4*nOut]
|
||
else if (b) {
|
||
std::vector<sd::LongType> dims = {0};
|
||
NDArray *dLdzReduced = dLdz.reduceAlongDimension(reduce::Sum, &dims); // [bS, 4*nOut] -> reduce -> [4*nOut]
|
||
*dLdb += *dLdzReduced;
|
||
delete dLdzReduced;
|
||
}
|
||
|
||
// dLdWp
|
||
if (Wp && x->rankOf() == 1) {
|
||
NDArray *firstOut = (*dLdWp)({0, nOut});
|
||
NDArray *secondOut = (*dLdWp)({nOut, 2 * nOut});
|
||
NDArray *thirdOut = (*dLdWp)({2 * nOut, 3 * nOut});
|
||
|
||
NDArray *dLdziMulCI = (*dLdzi) * (*cI); // [nOut]
|
||
*firstOut += *dLdziMulCI;
|
||
delete dLdziMulCI;
|
||
|
||
NDArray *dLdzfMulCI = (*dLdzf) * (*cI); // [nOut]
|
||
*secondOut += *dLdzfMulCI;
|
||
delete dLdzfMulCI;
|
||
|
||
NDArray *dLdzoMulC = (*dLdzo) * (*c); // [nOut]
|
||
*thirdOut += *dLdzoMulC;
|
||
delete dLdzoMulC;
|
||
|
||
delete firstOut;
|
||
delete secondOut;
|
||
delete thirdOut;
|
||
} else if (Wp) {
|
||
std::vector<sd::LongType> shape = {nOut};
|
||
NDArray temp2(Wp->ordering(), shape, Wp->dataType(), Wp->getContext());
|
||
std::vector<sd::LongType> dims = {0};
|
||
|
||
NDArray *firstOut = (*dLdWp)({0, nOut});
|
||
NDArray *secondOut = (*dLdWp)({nOut, 2 * nOut});
|
||
NDArray *thirdOut = (*dLdWp)({2 * nOut, 3 * nOut});
|
||
|
||
NDArray *dLdziMulCI = (*dLdzi) * (*cI);
|
||
dLdziMulCI->reduceAlongDimension(reduce::Sum, &temp2, &dims); // [bS, nOut] -> reduce -> [nOut]
|
||
*firstOut += temp2;
|
||
delete dLdziMulCI;
|
||
|
||
NDArray *dLdzfMulCI = (*dLdzf) * (*cI);
|
||
dLdzfMulCI->reduceAlongDimension(reduce::Sum, &temp2, &dims); // [bS, nOut] -> reduce -> [nOut]
|
||
*secondOut += temp2;
|
||
delete dLdzfMulCI;
|
||
|
||
NDArray *dLdzoMulC = (*dLdzo) * (*c);
|
||
dLdzoMulC->reduceAlongDimension(reduce::Sum, &temp2, &dims); // [bS, nOut] -> reduce -> [nOut]
|
||
*thirdOut += temp2;
|
||
delete dLdzoMulC;
|
||
|
||
delete firstOut;
|
||
delete secondOut;
|
||
delete thirdOut;
|
||
}
|
||
|
||
delete zUlike;
|
||
delete cUlike;
|
||
delete dLdzoUlike;
|
||
delete dcdcI;
|
||
}
|
||
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void lstmLayerTimeLoop(NDArray* x, NDArray* Wx, NDArray* Wr, NDArray* b, NDArray* seqLen,
|
||
NDArray* hI, NDArray* cI, NDArray* Wp, const std::vector<float>& params,
|
||
const bool forward, NDArray* h, NDArray* hL, NDArray* cL) {
|
||
// INPUTS:
|
||
// x - current input [sL, bS, nIn], [bS, sL, nIn], [bS, nIn, sL],
|
||
// Wx - input weights [nIn, 4*nOut]
|
||
// Wr - recurrent weights [nOut, 4*nOut]
|
||
// b - biases [4*nOut], optional, may be nullptr
|
||
// seqLen - [bS], optional, may be nullptr
|
||
// hI - initial output [bS, nOut], optional, may be nullptr
|
||
// cI - initial cell state at time t-1 [bS, nOut], optional, may be nullptr
|
||
// Wp - peephole weights [3*nOut], optional, may be nullptr
|
||
|
||
// OUTPUTS:
|
||
// h - output [sL, bS, nOut], [bS, sL, nOut], [bS, nOut, sL], optional, may be nullptr
|
||
// hL - output at last step [bS, nOut], optional, may be nullptr
|
||
// cL - cell state at last step [bS, nOut], optional, may be nullptr
|
||
|
||
// params = {dataFormat, directionMode, cellClip, gateAct, gateAlpha, gateBeta, cellAct, cellAlpha, cellBeta, outAct,
|
||
// outAlpha, outBeta}; dataFormat: 0,3 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL]
|
||
|
||
const int dataFormat = params[0];
|
||
const int directionMode = params[1];
|
||
|
||
const sd::LongType sL = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
|
||
const sd::LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
|
||
const sd::LongType nOut = Wx->sizeAt(-1) / 4;
|
||
|
||
std::vector<sd::LongType> shapeOut = {bS, nOut};
|
||
|
||
const auto type = h ? h->dataType() : (hL ? hL->dataType() : cL->dataType());
|
||
|
||
auto h0 = const_cast<NDArray*>(hI);
|
||
if (!hI) {
|
||
h0 = new NDArray(x->ordering(), shapeOut, type, x->getContext());
|
||
h0->nullify();
|
||
}
|
||
|
||
auto c0 = const_cast<NDArray*>(cI);
|
||
if (!cI) {
|
||
c0 = new NDArray(x->ordering(), shapeOut, type, x->getContext());
|
||
c0->nullify();
|
||
}
|
||
|
||
auto ct = cL;
|
||
if (!cL) ct = new NDArray(x->ordering(), shapeOut, type, x->getContext());
|
||
|
||
auto ht = hL;
|
||
if (!h && !hL) ht = new NDArray(x->ordering(), shapeOut, type, x->getContext());
|
||
|
||
// create sets of required (depends on seqLen presence) sub-arrays
|
||
std::vector<sd::LongType> *dims;
|
||
ResultSet *xSet(nullptr), *hSet(nullptr), *h0Set(nullptr), *c0Set(nullptr), *htSet(nullptr), *ctSet(nullptr);
|
||
|
||
if (!seqLen) {
|
||
std::vector<sd::LongType> dims2 = {dataFormat < 3 ? dataFormat : 0};
|
||
dims = ShapeUtils::evalDimsToExclude(x->rankOf(),
|
||
dims2.size(),dims2.data()); // points on bS and nIn/nOut axes
|
||
|
||
xSet = new ResultSet(x->allTensorsAlongDimension(*dims)); // sub-arrays with shape [bS, nIn]
|
||
if (h) hSet = new ResultSet(h->allTensorsAlongDimension(*dims)); // sub-arrays with shape [bS, nOut]
|
||
} else {
|
||
dims = dataFormat == 2 ? new std::vector<sd::LongType>({1}) : new std::vector<sd::LongType>({2}); // points on nIn/nOut axis
|
||
|
||
xSet = new ResultSet(x->allTensorsAlongDimension(*dims)); // sub-arrays with shape [nIn]
|
||
h0Set = new ResultSet(h0->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
c0Set = new ResultSet(c0->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
ctSet = new ResultSet(ct->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
if (h) hSet = new ResultSet(h->allTensorsAlongDimension(*dims)); // sub-arrays with shape [nOut]
|
||
if (ht) htSet = new ResultSet(ht->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
|
||
delete dims;
|
||
}
|
||
|
||
// loops
|
||
if (forward) {
|
||
if (!seqLen) {
|
||
if (!h) { // seqLen and h are absent
|
||
|
||
lstmLayerCell(xSet->at(0), Wx, Wr, b, h0, c0, Wp, params, ht, ct); // first time step
|
||
for (sd::LongType t = 1; t < sL; ++t)
|
||
lstmLayerCell(xSet->at(t), Wx, Wr, b, ht, ct, Wp, params, ht, ct); // rest time steps
|
||
} else { // seqLen is absent and h is present
|
||
|
||
lstmLayerCell(xSet->at(0), Wx, Wr, b, h0, c0, Wp, params, hSet->at(0), ct); // first time step
|
||
for (sd::LongType t = 1; t < sL; ++t)
|
||
lstmLayerCell(xSet->at(t), Wx, Wr, b, hSet->at(t - 1), ct, Wp, params, hSet->at(t), ct); // rest time steps
|
||
|
||
if (hL) hL->assign(hSet->at(sL - 1)); // assign last output to hL if it is not nullptr
|
||
}
|
||
} else {
|
||
if (!h) { // seqLen is present and h is absent
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
const int limit = seqLen->e<int>(e);
|
||
|
||
if (limit == 0) {
|
||
if (cL) ctSet->at(e)->nullify();
|
||
if (hL) htSet->at(e)->nullify();
|
||
continue;
|
||
}
|
||
|
||
auto ind = getBatchTimeTotalIndex(dataFormat, sL, bS, 0, e);
|
||
lstmLayerCell(xSet->at(ind), Wx, Wr, b, h0Set->at(e), c0Set->at(e), Wp, params, htSet->at(e),
|
||
ctSet->at(e)); // first time step
|
||
|
||
for (int t = 1; t < limit; ++t) {
|
||
ind = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
lstmLayerCell(xSet->at(ind), Wx, Wr, b, htSet->at(e), ctSet->at(e), Wp, params, htSet->at(e),
|
||
ctSet->at(e)); // rest time steps
|
||
}
|
||
}
|
||
} else { // seqLen and h are present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
int limit = seqLen->e<int>(e);
|
||
|
||
if (limit == 0) {
|
||
tensorAlongTimeBatchDims(*h, dataFormat, 0, 0, e, e + 1)
|
||
->nullify(); // nullify for given e and whole time range
|
||
|
||
if (cL) ctSet->at(e)->nullify();
|
||
if (hL) htSet->at(e)->nullify();
|
||
|
||
continue;
|
||
}
|
||
|
||
auto indPrev = getBatchTimeTotalIndex(dataFormat, sL, bS, 0, e);
|
||
lstmLayerCell(xSet->at(indPrev), Wx, Wr, b, h0Set->at(e), c0Set->at(e), Wp, params, hSet->at(indPrev),
|
||
ctSet->at(e)); // first time step
|
||
|
||
for (int t = 1; t < limit; ++t) {
|
||
auto indCurr = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
lstmLayerCell(xSet->at(indCurr), Wx, Wr, b, hSet->at(indPrev), ctSet->at(e), Wp, params, hSet->at(indCurr),
|
||
ctSet->at(e)); // rest time steps
|
||
indPrev = indCurr;
|
||
}
|
||
|
||
if (hL) htSet->at(e)->assign(hSet->at(indPrev)); // assign last output to hL if hL is not nullptr
|
||
|
||
if (limit != sL)
|
||
tensorAlongTimeBatchDims(*h, dataFormat, limit, sL, e, e + 1)
|
||
->nullify(); // nullify for given e and time range [limit, sL)
|
||
}
|
||
}
|
||
}
|
||
} else { // backward
|
||
|
||
if (!seqLen) {
|
||
if (!h) { // seqLen and h are absent
|
||
|
||
lstmLayerCell(xSet->at(sL - 1), Wx, Wr, b, h0, c0, Wp, params, ht, ct); // first time step
|
||
for (sd::LongType t = sL - 2; t >= 0; --t)
|
||
lstmLayerCell(xSet->at(t), Wx, Wr, b, ht, ct, Wp, params, ht, ct); // rest time steps
|
||
} else { // seqLen is absent and h is present
|
||
|
||
lstmLayerCell(xSet->at(sL - 1), Wx, Wr, b, h0, c0, Wp, params, hSet->at(sL - 1), ct); // first time step
|
||
for (sd::LongType t = sL - 2; t >= 0; --t)
|
||
lstmLayerCell(xSet->at(t), Wx, Wr, b, hSet->at(t + 1), ct, Wp, params, hSet->at(t), ct); // rest time steps
|
||
|
||
if (hL) hL->assign(hSet->at(0)); // assign last output to hL if it is not nullptr
|
||
}
|
||
} else if (directionMode == 1) { // only backward, no bidirectional mode
|
||
|
||
if (!h) { // h is absent and seqLen is present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
const int limit = seqLen->e<int>(e);
|
||
|
||
if (limit == 0) {
|
||
if (cL) ctSet->at(e)->nullify();
|
||
if (hL) htSet->at(e)->nullify();
|
||
continue;
|
||
}
|
||
|
||
auto ind = getBatchTimeTotalIndex(dataFormat, sL, bS, sL - 1, e);
|
||
lstmLayerCell(xSet->at(ind), Wx, Wr, b, h0Set->at(e), c0Set->at(e), Wp, params, htSet->at(e),
|
||
ctSet->at(e)); // first time step
|
||
|
||
for (sd::LongType t = sL - 2; t >= sL - limit; --t) {
|
||
ind = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
lstmLayerCell(xSet->at(ind), Wx, Wr, b, htSet->at(e), ctSet->at(e), Wp, params, htSet->at(e),
|
||
ctSet->at(e)); // rest time steps
|
||
}
|
||
}
|
||
} else { // seqLen and h are present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
int limit = seqLen->e<int>(e);
|
||
|
||
if (limit == 0) {
|
||
tensorAlongTimeBatchDims(*h, dataFormat, 0, 0, e, e + 1)
|
||
->nullify(); // nullify for given e and whole time range
|
||
|
||
if (cL) ctSet->at(e)->nullify();
|
||
if (hL) htSet->at(e)->nullify();
|
||
|
||
continue;
|
||
}
|
||
|
||
auto indPrev = getBatchTimeTotalIndex(dataFormat, sL, bS, sL - 1, e);
|
||
lstmLayerCell(xSet->at(indPrev), Wx, Wr, b, h0Set->at(e), c0Set->at(e), Wp, params, hSet->at(indPrev),
|
||
ctSet->at(e)); // first time step
|
||
|
||
for (sd::LongType t = sL - 2; t >= sL - limit; --t) {
|
||
auto indCurr = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
lstmLayerCell(xSet->at(indCurr), Wx, Wr, b, hSet->at(indPrev), ctSet->at(e), Wp, params, hSet->at(indCurr),
|
||
ctSet->at(e)); // rest time steps
|
||
indPrev = indCurr;
|
||
}
|
||
|
||
if (hL) htSet->at(e)->assign(hSet->at(indPrev)); // assign last output to hL if it is not nullptr
|
||
|
||
if (limit != sL)
|
||
tensorAlongTimeBatchDims(*h, dataFormat, 0, sL - limit, e, e + 1)
|
||
->nullify(); // nullify for given e and time range [limit, sL)
|
||
}
|
||
}
|
||
} else { // backward in bidirectional mode
|
||
|
||
if (!h) { // h is absent and seqLen is present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
const int limit = seqLen->e<int>(e);
|
||
|
||
if (limit == 0) {
|
||
if (cL) ctSet->at(e)->nullify();
|
||
if (hL) htSet->at(e)->nullify();
|
||
continue;
|
||
}
|
||
|
||
auto ind = getBatchTimeTotalIndex(dataFormat, sL, bS, limit - 1, e);
|
||
lstmLayerCell(xSet->at(ind), Wx, Wr, b, h0Set->at(e), c0Set->at(e), Wp, params, htSet->at(e),
|
||
ctSet->at(e)); // first time step
|
||
|
||
for (int t = limit - 2; t >= 0; --t) {
|
||
ind = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
lstmLayerCell(xSet->at(ind), Wx, Wr, b, htSet->at(e), ctSet->at(e), Wp, params, htSet->at(e),
|
||
ctSet->at(e)); // rest time steps
|
||
}
|
||
}
|
||
} else { // seqLen and h are present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
int limit = seqLen->e<int>(e);
|
||
|
||
if (limit == 0) {
|
||
tensorAlongTimeBatchDims(*h, dataFormat, 0, 0, e, e + 1)
|
||
->nullify(); // nullify for given e and whole time range
|
||
|
||
if (cL) ctSet->at(e)->nullify();
|
||
if (hL) htSet->at(e)->nullify();
|
||
|
||
continue;
|
||
}
|
||
|
||
auto indPrev = getBatchTimeTotalIndex(dataFormat, sL, bS, limit - 1, e);
|
||
lstmLayerCell(xSet->at(indPrev), Wx, Wr, b, h0Set->at(e), c0Set->at(e), Wp, params, hSet->at(indPrev),
|
||
ctSet->at(e)); // first time step
|
||
|
||
for (int t = limit - 2; t >= 0; --t) {
|
||
auto indCurr = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
lstmLayerCell(xSet->at(indCurr), Wx, Wr, b, hSet->at(indPrev), ctSet->at(e), Wp, params, hSet->at(indCurr),
|
||
ctSet->at(e)); // rest time steps
|
||
indPrev = indCurr;
|
||
}
|
||
|
||
if (hL) htSet->at(e)->assign(hSet->at(indPrev)); // assign last output to hL if it is not nullptr
|
||
|
||
if (limit != sL)
|
||
tensorAlongTimeBatchDims(*h, dataFormat, limit, sL, e, e + 1)
|
||
->nullify(); // nullify for given e and time range [limit, sL)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
delete xSet;
|
||
delete hSet;
|
||
delete h0Set;
|
||
delete c0Set;
|
||
delete htSet;
|
||
delete ctSet;
|
||
|
||
if (!hI) delete h0;
|
||
if (!cI) delete c0;
|
||
if (!cL) delete ct;
|
||
if (!h && !hL) delete ht;
|
||
}
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
void lstmLayerTimeLoopBp(NDArray* x, NDArray* Wx, NDArray* Wr, NDArray* b, NDArray* seqLen, NDArray* hI, NDArray* cI,
|
||
NDArray* Wp, NDArray* dLdh, NDArray* dLdhL, NDArray* dLdcL,
|
||
const std::vector<float>& params, const bool forward, NDArray* dLdx, NDArray* dLdWx,
|
||
NDArray* dLdWr, NDArray* dLdb, NDArray* dLdhI, NDArray* dLdcI, NDArray* dLdWp) {
|
||
// INPUTS:
|
||
// x - current input [sL, bS, nIn], [bS, sL, nIn], [bS, nIn, sL],
|
||
// Wx - input weights [nIn, 4*nOut]
|
||
// Wr - recurrent weights [nOut, 4*nOut]
|
||
// b - biases [4*nOut], optional, may be nullptr
|
||
// seqLen - [bS], optional, may be nullptr
|
||
// hI - initial output [bS, nOut], optional, may be nullptr
|
||
// cI - initial cell state at time t-1 [bS, nOut], optional, may be nullptr
|
||
// Wp - peephole weights [3*nOut], optional, may be nullptr
|
||
// dLdh - gradient vs. output [sL, bS, nOut], [bS, sL, nOut], [bS, nOut, sL], optional, may be nullptr
|
||
// dLdhL - gradient vs. output at last time step [bS, nOut], optional, may be nullptr
|
||
// dLdcL - gradient vs. cell state at last time step [bS, nOut], optional, may be nullptr
|
||
|
||
// OUTPUTS:
|
||
// dLdx - gradient vs. input [sL, bS, nIn], [bS, sL, nIn], [bS, nIn, sL]
|
||
// dLdWx - gradient vs. input weights [nIn, 4*nOut]
|
||
// dLdWr - gradient vs. recurrent weights [nOut, 4*nOut]
|
||
// dLdb - gradient vs. biases [4*nOut], optional, may be nullptr
|
||
// dLdhI - gradient vs. initial output [bS, nOut], optional, may be nullptr
|
||
// dLdcI - gradient vs. initial cell state at time t-1 [bS, nOut], optional, may be nullptr
|
||
// dLdWp - gradient vs. peephole weights [3*nOut], optional, may be nullptr
|
||
|
||
// params = {dataFormat, directionMode, cellClip, gateAct, gateAlpha, gateBeta, cellAct, cellAlpha, cellBeta, outAct,
|
||
// outAlpha, outBeta}; dataFormat: 0,3 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL]
|
||
|
||
const int dataFormat = params[0];
|
||
const int directionMode = params[1];
|
||
|
||
const int sL = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
|
||
const int bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
|
||
const int nOut = Wx->sizeAt(-1) / 4;
|
||
|
||
const auto type = dLdh ? dLdh->dataType() : (dLdhL ? dLdhL->dataType() : dLdcL->dataType());
|
||
|
||
std::vector<sd::LongType> shape = {bS, nOut};
|
||
auto dLdh0 = dLdhI;
|
||
if (!hI)
|
||
dLdh0 = new NDArray(x->ordering(), shape, type,
|
||
x->getContext()); // this constructor nullifies array automatically
|
||
|
||
auto dLdc0 = dLdcI;
|
||
if (!cI)
|
||
dLdc0 = new NDArray(x->ordering(), shape, type,
|
||
x->getContext()); // this constructor nullifies array automatically
|
||
|
||
std::vector<LongType> zShape = {sL, bS, 4 * nOut};
|
||
NDArray z(x->ordering(), zShape, type, x->getContext());
|
||
NDArray *zUlike2 = z.ulike();
|
||
NDArray a = *zUlike2;
|
||
std::vector<LongType> hShape = {sL + 1, bS, nOut};
|
||
NDArray h(x->ordering(), hShape, type, x->getContext());
|
||
NDArray *hUlike = h.ulike();
|
||
NDArray c = *hUlike;
|
||
|
||
// create sets of required (depends on seqLen presence) sub-arrays
|
||
std::vector<sd::LongType> *dims;
|
||
ResultSet *xSet(nullptr), *dLdxSet(nullptr), *hSet(nullptr), *cSet(nullptr), *zSet(nullptr), *aSet(nullptr),
|
||
*dLdhSet(nullptr), *dLdh0Set(nullptr), *dLdc0Set(nullptr), *dLdhLSet(nullptr), *dLdcLSet(nullptr),
|
||
*hISet(nullptr), *cISet(nullptr);
|
||
|
||
if (!seqLen) {
|
||
std::vector<sd::LongType> dim = {dataFormat < 3 ? dataFormat : 0};
|
||
dims = ShapeUtils::evalDimsToExclude(x->rankOf(),dim.size(),dim.data()); // points on [bS, nIn/nOut]
|
||
xSet = new ResultSet(x->allTensorsAlongDimension(*dims)); // sub-arrays with shape [bS, nIn]
|
||
dLdxSet = new ResultSet(dLdx->allTensorsAlongDimension(*dims)); // sub-arrays with shape [bS, nIn]
|
||
hSet = new ResultSet(h.allTensorsAlongDimension({1, 2})); // sub-arrays with shape [bS, nOut]
|
||
cSet = new ResultSet(c.allTensorsAlongDimension({1, 2})); // sub-arrays with shape [bS, nOut]
|
||
zSet = new ResultSet(z.allTensorsAlongDimension({1, 2})); // sub-arrays with shape [bS, 4*nOut]
|
||
aSet = new ResultSet(a.allTensorsAlongDimension({1, 2})); // sub-arrays with shape [bS, 4*nOut]
|
||
if (dLdh) dLdhSet = new ResultSet(dLdh->allTensorsAlongDimension(*dims)); // sub-arrays with shape [bS, nOut]
|
||
|
||
} else {
|
||
dims = dataFormat == 2 ? new std::vector<sd::LongType>({1}) : new std::vector<sd::LongType>({2}); // points on nIn/nOut axis
|
||
|
||
xSet = new ResultSet(x->allTensorsAlongDimension(*dims)); // sub-arrays with shape [nIn]
|
||
dLdxSet = new ResultSet(dLdx->allTensorsAlongDimension(*dims)); // sub-arrays with shape [nIn]
|
||
hSet = new ResultSet(h.allTensorsAlongDimension({2})); // sub-arrays with shape [nOut]
|
||
cSet = new ResultSet(c.allTensorsAlongDimension({2})); // sub-arrays with shape [nOut]
|
||
zSet = new ResultSet(z.allTensorsAlongDimension({2})); // sub-arrays with shape [4*nOut]
|
||
aSet = new ResultSet(a.allTensorsAlongDimension({2})); // sub-arrays with shape [4*nOut]
|
||
|
||
if (hI) hISet = new ResultSet(hI->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
if (cI) cISet = new ResultSet(cI->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
|
||
dLdh0Set = new ResultSet(dLdh0->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
dLdc0Set = new ResultSet(dLdc0->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
|
||
if (dLdh) dLdhSet = new ResultSet(dLdh->allTensorsAlongDimension(*dims)); // sub-arrays with shape [nOut]
|
||
if (dLdhL) dLdhLSet = new ResultSet(dLdhL->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
if (dLdcL) dLdcLSet = new ResultSet(dLdcL->allTensorsAlongDimension({1})); // sub-arrays with shape [nOut]
|
||
|
||
delete dims;
|
||
}
|
||
|
||
// loops
|
||
if (forward) {
|
||
if (!seqLen) { // seqLen is absent
|
||
|
||
if (hI)
|
||
hSet->at(0)->assign(hI);
|
||
else
|
||
hSet->at(0)->nullify();
|
||
if (cI)
|
||
cSet->at(0)->assign(cI);
|
||
else
|
||
cSet->at(0)->nullify();
|
||
|
||
// ff
|
||
for (sd::LongType t = 0; t < sL; ++t) {
|
||
lstmLayerCell(xSet->at(t), Wx, Wr, b, hSet->at(t), cSet->at(t), Wp, params, zSet->at(t), aSet->at(t),
|
||
hSet->at(t + 1), cSet->at(t + 1));
|
||
}
|
||
|
||
// bp
|
||
for (sd::LongType t = sL - 1; t >= 0; --t) {
|
||
|
||
NDArray* dLdhh = dLdh ? dLdhSet->at(t) : nullptr;
|
||
NDArray* dLdhhL = (t == sL - 1 && dLdhL) ? dLdhL : nullptr;
|
||
NDArray* dLdccL = (t == sL - 1 && dLdcL) ? dLdcL : nullptr;
|
||
lstmLayerCellBp(xSet->at(t), Wx, Wr, b, hSet->at(t), cSet->at(t), Wp, dLdhh, dLdhhL, dLdccL, zSet->at(t),
|
||
aSet->at(t), cSet->at(t + 1), params, dLdxSet->at(t), dLdWx, dLdWr, dLdh0, dLdc0, dLdb, dLdWp);
|
||
|
||
}
|
||
|
||
|
||
|
||
} else { // seqLen is present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
const sd::LongType limit = seqLen->e<sd::LongType>(e);
|
||
|
||
if (limit == 0) {
|
||
tensorAlongTimeBatchDims(*dLdx, dataFormat, 0, 0, e, e + 1)
|
||
->nullify(); // nullify for given e and whole time range
|
||
continue;
|
||
}
|
||
|
||
if (hI)
|
||
hSet->at(e)->assign(hISet->at(e));
|
||
else
|
||
hSet->at(e)->nullify();
|
||
if (cI)
|
||
cSet->at(e)->assign(cISet->at(e));
|
||
else
|
||
cSet->at(e)->nullify();
|
||
|
||
// ff
|
||
for (sd::LongType t = 0; t < limit; ++t) {
|
||
lstmLayerCell(xSet->at(getBatchTimeTotalIndex(dataFormat, sL, bS, t, e)), Wx, Wr, b, hSet->at(t * bS + e),
|
||
cSet->at(t * bS + e), Wp, params, zSet->at(t * bS + e), aSet->at(t * bS + e),
|
||
hSet->at((t + 1) * bS + e), cSet->at((t + 1) * bS + e));
|
||
}
|
||
// bp
|
||
for (sd::LongType t = limit - 1; t >= 0; --t) {
|
||
const auto ind = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
NDArray* dLdhh = dLdh ? dLdhSet->at(ind) : nullptr;
|
||
NDArray* dLdhhL = (t == limit - 1 && dLdhL) ? dLdhLSet->at(e) : nullptr;
|
||
NDArray* dLdccL = (t == limit - 1 && dLdcL) ? dLdcLSet->at(e) : nullptr;
|
||
lstmLayerCellBp(xSet->at(ind), Wx, Wr, b, hSet->at(t * bS + e), cSet->at(t * bS + e), Wp, dLdhh, dLdhhL,
|
||
dLdccL, zSet->at(t * bS + e), aSet->at(t * bS + e), cSet->at((t + 1) * bS + e), params,
|
||
dLdxSet->at(ind), dLdWx, dLdWr, dLdh0Set->at(e), dLdc0Set->at(e), dLdb, dLdWp);
|
||
}
|
||
|
||
if (limit != sL)
|
||
tensorAlongTimeBatchDims(*dLdx, dataFormat, limit, sL, e, e + 1)
|
||
->nullify(); // nullify for given e and time range [limit, sL)
|
||
}
|
||
|
||
}
|
||
} else { // backward or bidirectional
|
||
|
||
if (!seqLen) { // backward or bidirectional, seqLen is absent
|
||
|
||
if (hI)
|
||
hSet->at(sL)->assign(hI);
|
||
else
|
||
hSet->at(sL)->nullify();
|
||
if (cI)
|
||
cSet->at(sL)->assign(cI);
|
||
else
|
||
cSet->at(sL)->nullify();
|
||
|
||
// ff
|
||
for (sd::LongType t = sL - 1; t >= 0; --t) {
|
||
lstmLayerCell(xSet->at(t), Wx, Wr, b, hSet->at(t + 1), cSet->at(t + 1), Wp, params, zSet->at(t), aSet->at(t),
|
||
hSet->at(t), cSet->at(t));
|
||
}
|
||
|
||
// bp
|
||
for (sd::LongType t = 0; t < sL; ++t) {
|
||
NDArray* dLdhh = dLdh ? dLdhSet->at(t) : nullptr;
|
||
NDArray* dLdhhL = (t == 0 && dLdhL) ? dLdhL : nullptr;
|
||
NDArray* dLdccL = (t == 0 && dLdcL) ? dLdcL : nullptr;
|
||
lstmLayerCellBp(xSet->at(t), Wx, Wr, b, hSet->at(t + 1), cSet->at(t + 1), Wp, dLdhh, dLdhhL, dLdccL,
|
||
zSet->at(t), aSet->at(t), cSet->at(t), params, dLdxSet->at(t), dLdWx, dLdWr, dLdh0, dLdc0, dLdb,
|
||
dLdWp);
|
||
}
|
||
|
||
|
||
|
||
} else if (directionMode == 1) { // backward, seqLen is present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
const sd::LongType limit = seqLen->e<sd::LongType>(e);
|
||
|
||
if (limit == 0) {
|
||
tensorAlongTimeBatchDims(*dLdx, dataFormat, 0, 0, e, e + 1)
|
||
->nullify(); // nullify for given e and whole time range
|
||
continue;
|
||
}
|
||
|
||
if (hI)
|
||
hSet->at(sL * bS + e)->assign(hISet->at(e));
|
||
else
|
||
hSet->at(sL * bS + e)->nullify();
|
||
if (cI)
|
||
cSet->at(sL * bS + e)->assign(cISet->at(e));
|
||
else
|
||
cSet->at(sL * bS + e)->nullify();
|
||
|
||
// ff
|
||
for (int t = sL - 1; t >= sL - limit; --t)
|
||
lstmLayerCell(xSet->at(getBatchTimeTotalIndex(dataFormat, sL, bS, t, e)), Wx, Wr, b,
|
||
hSet->at((t + 1) * bS + e), cSet->at((t + 1) * bS + e), Wp, params, zSet->at(t * bS + e),
|
||
aSet->at(t * bS + e), hSet->at(t * bS + e), cSet->at(t * bS + e));
|
||
|
||
// bp
|
||
for (sd::LongType t = sL - limit; t < sL; ++t) {
|
||
const auto ind = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
NDArray* dLdhh = dLdh ? dLdhSet->at(ind) : nullptr;
|
||
NDArray* dLdhhL = (t == sL - limit && dLdhL) ? dLdhLSet->at(e) : nullptr;
|
||
NDArray* dLdccL = (t == sL - limit && dLdcL) ? dLdcLSet->at(e) : nullptr;
|
||
lstmLayerCellBp(xSet->at(ind), Wx, Wr, b, hSet->at((t + 1) * bS + e), cSet->at((t + 1) * bS + e), Wp, dLdhh,
|
||
dLdhhL, dLdccL, zSet->at(t * bS + e), aSet->at(t * bS + e), cSet->at(t * bS + e), params,
|
||
dLdxSet->at(ind), dLdWx, dLdWr, dLdh0Set->at(e), dLdc0Set->at(e), dLdb, dLdWp);
|
||
}
|
||
|
||
if (limit != sL)
|
||
tensorAlongTimeBatchDims(*dLdx, dataFormat, 0, sL - limit, e, e + 1)
|
||
->nullify(); // nullify for given e and time range [limit, sL)
|
||
}
|
||
|
||
|
||
} else { // bidirectional mode, seqLen is present
|
||
|
||
for (sd::LongType e = 0; e < bS; ++e) {
|
||
const int limit = seqLen->e<sd::LongType>(e);
|
||
|
||
if (limit == 0) {
|
||
tensorAlongTimeBatchDims(*dLdx, dataFormat, 0, 0, e, e + 1)
|
||
->nullify(); // nullify for given e and whole time range
|
||
continue;
|
||
}
|
||
|
||
NDArray *hView = h({limit, limit + 1, e, e + 1, 0, 0});
|
||
NDArray *cView = c({limit, limit + 1, e, e + 1, 0, 0});
|
||
|
||
if (hI)
|
||
hView->assign(hISet->at(e));
|
||
else
|
||
hView->nullify();
|
||
if (cI)
|
||
cView->assign(cISet->at(e));
|
||
else
|
||
cView->nullify();
|
||
|
||
delete hView;
|
||
delete cView;
|
||
|
||
// ff
|
||
for (int t = limit - 1; t >= 0; --t)
|
||
lstmLayerCell(xSet->at(getBatchTimeTotalIndex(dataFormat, sL, bS, t, e)), Wx, Wr, b,
|
||
hSet->at((t + 1) * bS + e), cSet->at((t + 1) * bS + e), Wp, params, zSet->at(t * bS + e),
|
||
aSet->at(t * bS + e), hSet->at(t * bS + e), cSet->at(t * bS + e));
|
||
|
||
// bp
|
||
for (sd::LongType t = 0; t < limit; ++t) {
|
||
const auto ind = getBatchTimeTotalIndex(dataFormat, sL, bS, t, e);
|
||
NDArray* dLdhh = dLdh ? dLdhSet->at(ind) : nullptr;
|
||
NDArray* dLdhhL = (t == 0 && dLdhL) ? dLdhLSet->at(e) : nullptr;
|
||
NDArray* dLdccL = (t == 0 && dLdcL) ? dLdcLSet->at(e) : nullptr;
|
||
lstmLayerCellBp(xSet->at(ind), Wx, Wr, b, hSet->at((t + 1) * bS + e), cSet->at((t + 1) * bS + e), Wp, dLdhh,
|
||
dLdhhL, dLdccL, zSet->at(t * bS + e), aSet->at(t * bS + e), cSet->at(t * bS + e), params,
|
||
dLdxSet->at(ind), dLdWx, dLdWr, dLdh0Set->at(e), dLdc0Set->at(e), dLdb, dLdWp);
|
||
}
|
||
|
||
if (limit != sL)
|
||
tensorAlongTimeBatchDims(*dLdx, dataFormat, limit, sL, e, e + 1)
|
||
->nullify(); // nullify for given e and time range [limit, sL)
|
||
}
|
||
|
||
|
||
}
|
||
}
|
||
|
||
delete xSet;
|
||
delete dLdxSet;
|
||
delete hSet;
|
||
delete cSet;
|
||
delete aSet;
|
||
delete zSet;
|
||
delete dLdhSet;
|
||
delete dLdh0Set;
|
||
delete dLdc0Set;
|
||
delete dLdhLSet;
|
||
delete dLdcLSet;
|
||
delete hISet;
|
||
delete cISet;
|
||
delete zUlike2;
|
||
delete hUlike;
|
||
if (!hI) delete dLdh0;
|
||
if (!cI) delete dLdc0;
|
||
}
|
||
|
||
} // namespace helpers
|
||
} // namespace ops
|
||
} // namespace sd
|
||
|
||
//////////////////////////////////////////////////////////////////////////
|
||
|
||
#endif
|