wehub-resource-sync
386 lines
14 KiB
C++
386 lines
14 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
|
|
******************************************************************************/
|
|
|
|
//
|
|
// implementation of operations for Simple Recurrent Unit: arXiv:1709.02755v2 [cs.CL] 12 Sep 2017
|
|
//
|
|
// @author Yurii Shyrma, created on 05.12.2017
|
|
//
|
|
#include <array/NDArrayFactory.h>
|
|
#include <execution/Threads.h>
|
|
#include <helpers/MmulHelper.h>
|
|
#include <ops/declarable/helpers/sru.h>
|
|
#if NOT_EXCLUDED(OP_sru)
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
static SD_INLINE NDArray activation(NDArray& arr) {
|
|
auto result = NDArray(&arr, false, arr.getContext());
|
|
(const_cast<NDArray&>(arr)).applyTransform(transform::Tanh, &result);
|
|
return result;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
static SD_INLINE NDArray* sigmoid(NDArray& arr) {
|
|
return (const_cast<NDArray&>(arr)).transform(transform::Sigmoid);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void sruCell(sd::LaunchContext* context, NDArray* x, NDArray* c0, NDArray* w, NDArray* b,
|
|
NDArray* h, NDArray* c) {
|
|
// x input [bS x inSize], bS - batch size, inSize - number of features
|
|
// c0 previous cell state c [bS x inSize], that is at previous time step t-1
|
|
// w weights [inSize x 3*inSize]
|
|
// b biases [2*inSize]
|
|
|
|
// h current cell output [bS x inSize], that is at current time step t
|
|
// c current cell state [bS x inSize], that is at current time step t
|
|
|
|
const int inSize = x->sizeAt(1); // inSize - number of features
|
|
|
|
NDArray *z = mmul(*x, *w); // [bS x 3*inSize]
|
|
|
|
// forget gate = sigmoid(x*Wf + bf)
|
|
NDArray *zView1 = (*z)({0, 0, inSize, 2 * inSize});
|
|
NDArray *bView1 = (*b)({0, inSize});
|
|
NDArray *addResult1 = (*zView1) + (*bView1);
|
|
NDArray *f = sigmoid(*addResult1);
|
|
delete addResult1;
|
|
|
|
// reset gate = sigmoid(x*Wr + br)
|
|
NDArray *zView2 = (*z)({0, 0, 2 * inSize, 3 * inSize});
|
|
NDArray *bView2 = (*b)({inSize, 2 * inSize});
|
|
NDArray *addResult2 = (*zView2) + (*bView2);
|
|
NDArray *r = sigmoid(*addResult2);
|
|
delete addResult2;
|
|
|
|
// ◦ means element-wise product or so called Hadamard product
|
|
// current sell state = f◦c0 + (1 - f)◦(x*Wc)
|
|
NDArray *zView3 = (*z)({0, 0, 0, inSize});
|
|
NDArray *fMulC0 = (*f) * (*c0);
|
|
NDArray *oneMinusF = 1.f - (*f);
|
|
NDArray *oneMinusFMulZ = (*oneMinusF) * (*zView3);
|
|
NDArray *assignOne = (*fMulC0) + (*oneMinusFMulZ);
|
|
c->assign(assignOne);
|
|
delete fMulC0;
|
|
delete oneMinusFMulZ;
|
|
delete assignOne;
|
|
delete oneMinusF;
|
|
// *c = f*(*c0 - z({},{0, inSize})) + z({{},{0, inSize}});
|
|
|
|
// current cell output = r◦activation(c) + (1 - r)◦x
|
|
NDArray activationC = activation(*c);
|
|
NDArray *rMulActivation = (*r) * activationC;
|
|
NDArray *oneMinusR = 1.f - (*r);
|
|
NDArray *oneMinusRMulX = *oneMinusR * (*x);
|
|
NDArray *assign2 = (*rMulActivation) + (*oneMinusRMulX);
|
|
h->assign(assign2);
|
|
delete rMulActivation;
|
|
delete oneMinusRMulX;
|
|
delete assign2;
|
|
delete oneMinusR;
|
|
// *h = r * (activation<T>(c) - *x) + *x;
|
|
|
|
delete z;
|
|
delete zView1;
|
|
delete bView1;
|
|
delete f;
|
|
delete zView2;
|
|
delete bView2;
|
|
delete r;
|
|
delete zView3;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void sruTimeLoop(sd::LaunchContext* context, NDArray* x, NDArray* c0, NDArray* w, NDArray* b,
|
|
NDArray* h, NDArray* c) {
|
|
// x input [bS x inSize x time]
|
|
// c0 initial cell state (at time step = 0) [bS x inSize],
|
|
// w weights, [3*inSize x inSize]
|
|
// b biases, [2*inSize]
|
|
|
|
// h cell outputs [bS x inSize x time]
|
|
// c cell states [bS x inSize x time]
|
|
|
|
NDArray *wT = w->transpose(); // [3*inSize x inSize] -> [inSize x 3*inSize]
|
|
|
|
const int time = x->sizeAt(2);
|
|
|
|
NDArray ct_1(*c0);
|
|
|
|
// loop through time steps
|
|
for (int t = 0; t < time; ++t) {
|
|
NDArray *xt = (*x)({0, 0, 0, 0, t, t + 1});
|
|
NDArray *ht = (*h)({0, 0, 0, 0, t, t + 1});
|
|
NDArray *ct = (*c)({0, 0, 0, 0, t, t + 1});
|
|
|
|
helpers::sruCell(context, xt, &ct_1, wT, b, ht, ct);
|
|
ct_1.assign(ct);
|
|
|
|
delete xt;
|
|
delete ht;
|
|
delete ct;
|
|
}
|
|
|
|
delete wT;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void sruBI_(NDArray* x, NDArray* w, NDArray* b, NDArray* c0, NDArray* mask, NDArray* ht,
|
|
NDArray* ct) {
|
|
// x input 3d tensor [time x bS x 2*K], time - number of time steps, bS - batch size, K - number of features
|
|
// w 2d tensor of weights [2*K x 6*K]
|
|
// b row of biases with twice length [4*K]
|
|
// c0 2d tensor of initial state [bS x 2*K] at time t=0
|
|
// mask optional, 2d tensor of dropout mask [bS x 2*K]
|
|
|
|
// ht [time x bS x 2*K]
|
|
// ct [time x bS x 2*K]
|
|
|
|
const sd::LongType time = x->sizeAt(0); // time - number of time steps
|
|
const sd::LongType bS = x->sizeAt(1); // bS - batch size
|
|
const sd::LongType K = x->sizeAt(2) / 2; // K - number of features
|
|
std::vector<sd::LongType> dims2 = {1, 2};
|
|
// x = x * mask
|
|
if (mask) x->applyBroadcast(broadcast::Multiply, &dims2, mask, x); // apply mask
|
|
|
|
// U = x * w
|
|
NDArray *wi = mmul(*x, *w); // U [time x bS x 6*K]
|
|
|
|
const sd::LongType d2 = 2 * K;
|
|
const sd::LongType ncols = bS * d2;
|
|
const sd::LongType ncolsWi = 3 * ncols;
|
|
|
|
T* pI = x->bufferAsT<T>();
|
|
T* pWi = wi->bufferAsT<T>();
|
|
T* pBias = const_cast<NDArray*>(b)->bufferAsT<T>();
|
|
T* pInit = const_cast<NDArray*>(c0)->bufferAsT<T>();
|
|
T* pMask = mask ? const_cast<NDArray*>(mask)->bufferAsT<T>() : nullptr;
|
|
T* pHt = ht->bufferAsT<T>();
|
|
T* pCt = ct->bufferAsT<T>();
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto col = start; col < stop; col++) {
|
|
const auto colNum = col % d2;
|
|
bool flip = colNum >= K;
|
|
T maskVal = mask ? *(pMask + col) : T(1);
|
|
T cur = *(pInit + col);
|
|
T bF = *(pBias + colNum);
|
|
T bR = *(pBias + colNum + d2);
|
|
T* pWiVal = pWi + 3 * col;
|
|
T* pIVal = pI + col;
|
|
T* pHtVal = pHt + col;
|
|
T* pCtVal = pCt + col;
|
|
|
|
if (flip) {
|
|
const auto step = (time - 1) * ncols;
|
|
pIVal += step;
|
|
pHtVal += step;
|
|
pCtVal += step;
|
|
pWiVal += (time - 1) * ncolsWi;
|
|
}
|
|
|
|
auto ncolsRev = flip ? -ncols : ncols;
|
|
auto ncolsWiRev = flip ? -ncolsWi : ncolsWi;
|
|
|
|
for (sd::LongType t = 0; t < time; ++t) {
|
|
// evaluate sigmoids
|
|
T ft = (1.) / (1. + sd::math::sd_exp<T, T>(-(pWiVal[1] + bF)));
|
|
T rt = (1.) / (1. + sd::math::sd_exp<T, T>(-(pWiVal[2] + bR)));
|
|
|
|
cur = (cur - *pWiVal) * ft + *pWiVal;
|
|
*pCtVal = cur;
|
|
T val = sd::math::sd_tanh<T, T>(cur);
|
|
*pHtVal = (val * maskVal - *pIVal) * rt + *pIVal;
|
|
|
|
pIVal += ncolsRev;
|
|
pWiVal += ncolsWiRev;
|
|
pCtVal += ncolsRev;
|
|
pHtVal += ncolsRev;
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, ncols);
|
|
|
|
delete wi;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void sruBIBP_(NDArray* x, NDArray* w, NDArray* b, NDArray* c0, NDArray* ct,
|
|
NDArray* inGradC0, NDArray* inGradHt, NDArray* mask, NDArray* gradI,
|
|
NDArray* gradW, NDArray* gradB, NDArray* gradC0) {
|
|
// x input 3d tensor [time x bS x 2*K], time - number of time steps, bS - batch size, K - number of features
|
|
// w 2d tensor of weights [2*K x 6*K]
|
|
// b row of biases with twice length 4*K]
|
|
// c0 2d tensor of initial state [bS x 2*K] at time t=0
|
|
// ct [time x bS x 2*K]
|
|
// inGradC0 [bS x 2*K]
|
|
// inGradHt [time x bS x 2*K]
|
|
// mask optional, 2d tensor of dropout mask [bS x 2*K]
|
|
|
|
// gradI [time x bS x 2*K]
|
|
// gradW [time x 2*K x 6*K]
|
|
// gradB [4*K]
|
|
// gradC0 [bS x 2*K]
|
|
|
|
const sd::LongType time = x->sizeAt(0); // time - number of time steps
|
|
const sd::LongType bS = x->sizeAt(1);
|
|
const sd::LongType K = x->sizeAt(2) / 2;
|
|
std::vector<sd::LongType> dims2 = {1, 2};
|
|
|
|
// x = x * mask
|
|
if (mask) x->applyBroadcast(broadcast::Multiply, &dims2, mask, x); // apply mask
|
|
|
|
// U = x * w
|
|
NDArray *wi = mmul(*x, *w); // [time x bS x 2*K] * [2*K x 6*K] = [time x bS x 6*K]
|
|
std::vector<sd::LongType> biasShape = {bS, 4 * K};
|
|
std::vector<sd::LongType> wShape = {time, bS, 6 * K};
|
|
NDArray gradBias(x->ordering(), biasShape, x->dataType(), x->getContext());
|
|
NDArray gradWi(x->ordering(), wShape, x->dataType(), x->getContext());
|
|
|
|
const sd::LongType d2 = 2 * K;
|
|
const sd::LongType ncols = bS * d2;
|
|
const sd::LongType ncolsWi = 3 * ncols;
|
|
T* pInput = x->bufferAsT<T>();
|
|
T* pWi = wi->bufferAsT<T>();
|
|
T* pBias = const_cast<NDArray*>(b)->bufferAsT<T>();
|
|
T* pInit = const_cast<NDArray*>(c0)->bufferAsT<T>();
|
|
T* pMask = mask ? const_cast<NDArray*>(mask)->bufferAsT<T>() : nullptr;
|
|
T* pState = const_cast<NDArray*>(ct)->bufferAsT<T>();
|
|
T* pInGradCt = const_cast<NDArray*>(inGradC0)->bufferAsT<T>();
|
|
T* pInGradHt = const_cast<NDArray*>(inGradHt)->bufferAsT<T>();
|
|
T* pGradWi = gradWi.bufferAsT<T>();
|
|
T* pGradInput = gradI->bufferAsT<T>();
|
|
T* pGradBias = gradBias.bufferAsT<T>();
|
|
T* pGradInit = gradC0->bufferAsT<T>();
|
|
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
for (auto col = start; col < stop; col++) {
|
|
T gbF = static_cast<T>(0.f);
|
|
T gbR = static_cast<T>(0.f);
|
|
const auto colNum = col % d2;
|
|
const bool flip = colNum >= K;
|
|
T maskVal = mask ? *(pMask + col) : T(1.);
|
|
T cur = *(pInGradCt + col);
|
|
T bF = *(pBias + colNum);
|
|
T bR = *(pBias + colNum + d2);
|
|
T* pWiVal = pWi + 3 * col;
|
|
T* pInputVal = pInput + col;
|
|
T* pStateVal = pState + col;
|
|
T* pInGradHtVal = pInGradHt + col;
|
|
T* pGradWiVal = pGradWi + 3 * col;
|
|
T* pGradInputVal = pGradInput + col;
|
|
|
|
if (!flip) {
|
|
const auto stepI = (time - 1) * ncols;
|
|
const auto stepW = (time - 1) * ncolsWi;
|
|
pInputVal += stepI;
|
|
pStateVal += stepI;
|
|
pInGradHtVal += stepI;
|
|
pGradInputVal += stepI;
|
|
pWiVal += stepW;
|
|
pGradWiVal += stepW;
|
|
}
|
|
|
|
sd::LongType ncolsRev = flip ? -ncols : ncols;
|
|
sd::LongType ncolsWiRev = flip ? -ncolsWi : ncolsWi;
|
|
|
|
for (sd::LongType t = 0; t < time; ++t) {
|
|
// evaluate sigmoids
|
|
T ft = ((T)1.) / ((T)1. + sd::math::sd_exp<T, T>(-(*(pWiVal + 1) + bF)));
|
|
T rt = ((T)1.) / ((T)1. + sd::math::sd_exp<T, T>(-(*(pWiVal + 2) + bR)));
|
|
|
|
T val = sd::math::sd_tanh<T, T>(*pStateVal);
|
|
T prevVal = (t < time - 1) ? (*(pStateVal - ncolsRev)) : (*(pInit + col));
|
|
// grad wrt input
|
|
*pGradInputVal = *pInGradHtVal - (*pInGradHtVal) * rt;
|
|
// grad wrt rt, wiR and bR
|
|
T grt = (*pInGradHtVal) * (val * maskVal - *pInputVal) * (rt - rt * rt);
|
|
*(pGradWiVal + 2) = grt;
|
|
gbR += grt;
|
|
// grad wrt state
|
|
T gradSateVal = (*pInGradHtVal) * maskVal * (rt - rt * val * val) + cur;
|
|
// grad wrt wi0
|
|
*pGradWiVal = gradSateVal - gradSateVal * ft;
|
|
// grad wrt ft, wi1, and bF
|
|
T gft = gradSateVal * (prevVal - *pWiVal) * (ft - ft * ft);
|
|
*(pGradWiVal + 1) = gft;
|
|
gbF += gft;
|
|
// grad wrt c_previous
|
|
cur = gradSateVal * ft;
|
|
|
|
pInputVal -= ncolsRev;
|
|
pWiVal -= ncolsWiRev;
|
|
pStateVal -= ncolsRev;
|
|
pGradWiVal -= ncolsWiRev;
|
|
pGradInputVal -= ncolsRev;
|
|
pInGradHtVal -= ncolsRev;
|
|
}
|
|
*(pGradBias + col) = gbF;
|
|
*(pGradBias + col + ncols) = gbR;
|
|
*(pGradInit + col) = cur;
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, ncols);
|
|
|
|
// gradB
|
|
|
|
std::vector<sd::LongType> dims = {0};
|
|
gradBias.reduceAlongDimension(reduce::Sum, gradB, &dims); // [4*K]
|
|
|
|
// gradW
|
|
x->permutei({0, 2, 1}, 0, false); // [time x bS x 2*K] -> [time x 2*K x bS]
|
|
MmulHelper::mmul(x, &gradWi, gradW, 1., 0.); // [time x 2*K x bS ] * [time x bS x 6*K] = [time x 2*K x 6*K]
|
|
|
|
delete wi;
|
|
}
|
|
|
|
void sruBI(sd::LaunchContext* context, NDArray* x, NDArray* w, NDArray* b, NDArray* c0,
|
|
NDArray* mask, NDArray* ht, NDArray* ct) {
|
|
BUILD_SINGLE_SELECTOR(x->dataType(), sruBI_, (x, w, b, c0, mask, ht, ct), SD_FLOAT_TYPES);
|
|
}
|
|
void sruBIBP(sd::LaunchContext* context, NDArray* x, NDArray* w, NDArray* b, NDArray* c0,
|
|
NDArray* ct, NDArray* inGradC0, NDArray* inGradH, NDArray* mask, NDArray* gradI,
|
|
NDArray* gradW, NDArray* gradB, NDArray* gradC0) {
|
|
BUILD_SINGLE_SELECTOR(x->dataType(), sruBIBP_,
|
|
(x, w, b, c0, ct, inGradC0, inGradH, mask, gradI, gradW, gradB, gradC0), SD_FLOAT_TYPES);
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( void sruBI_,
|
|
(NDArray * x, NDArray* w, NDArray* b, NDArray* c0, NDArray* mask,
|
|
NDArray* ht, NDArray* ct),
|
|
SD_FLOAT_TYPES);
|
|
BUILD_SINGLE_TEMPLATE( void sruBIBP_,
|
|
(NDArray * x, NDArray* w, NDArray* b, NDArray* c0, NDArray* ct,
|
|
NDArray* inGradC0, NDArray* inGradH, NDArray* mask, NDArray* gradI,
|
|
NDArray* gradW, NDArray* gradB, NDArray* gradC0),
|
|
SD_FLOAT_TYPES);
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
|
|
#endif
|