wehub-resource-sync
664 lines
34 KiB
C++
664 lines
34 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
|
|
//
|
|
|
|
#include <system/op_boilerplate.h>
|
|
#if NOT_EXCLUDED(OP_sru)
|
|
|
|
#include <helpers/MmulHelper.h>
|
|
#include <helpers/PointersManager.h>
|
|
#include <ops/declarable/CustomOperations.h>
|
|
#include <ops/declarable/helpers/sru.h>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
CUSTOM_OP_IMPL(sru, 5, 2, false, 0, 0) {
|
|
auto x = INPUT_VARIABLE(0); // X, input 3d tensor [bS x inSize x time], time - number of time steps, bS - batch size,
|
|
// inSize - number of features
|
|
auto w = INPUT_VARIABLE(1); // W, 2d tensor of weights [3*inSize x inSize]
|
|
auto b = INPUT_VARIABLE(2); // B, row of biases with twice length [2*inSize]
|
|
auto c0 = INPUT_VARIABLE(3); // C_{0}, 2d tensor of initial state [bS x inSize] at time t=0
|
|
auto mask = block.width() > 4 ? INPUT_VARIABLE(4) : nullptr; // optional, 2d tensor of dropout mask [bS x inSize]
|
|
|
|
auto h = OUTPUT_VARIABLE(0); // cell outputs, [bS x inSize x time]
|
|
auto c = OUTPUT_VARIABLE(1); // cell states, [bS x inSize x time]
|
|
|
|
const int rank = x->rankOf(); // = 3
|
|
const auto bS = x->sizeAt(0);
|
|
const auto inSize = x->sizeAt(1);
|
|
const auto time = x->sizeAt(2);
|
|
|
|
// input shapes validation
|
|
REQUIRE_TRUE(w->rankOf() == rank - 1, 0,
|
|
"SRU operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
|
|
w->rankOf());
|
|
REQUIRE_TRUE(b->rankOf() == 1, 0, "SRU operation: wrong rank of biases array, expected is %i, but got %i instead !",
|
|
1, b->rankOf());
|
|
REQUIRE_TRUE(c0->rankOf() == rank - 1, 0,
|
|
"SRU operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
|
|
c0->rankOf());
|
|
if (mask)
|
|
REQUIRE_TRUE(mask->rankOf() == rank - 1, 0,
|
|
"SRU operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
|
|
mask->rankOf());
|
|
|
|
const std::vector<LongType> wCorrectShape = {3 * inSize, inSize};
|
|
const std::vector<LongType> bCorrectShape = {2 * inSize};
|
|
const std::vector<LongType> c0CorrectShape = {bS, inSize};
|
|
|
|
REQUIRE_TRUE(w->isSameShape(wCorrectShape), 0,
|
|
"SRU operation: wrong shape of weights array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
|
|
REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
|
|
"SRU operation: wrong shape of biases array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
|
|
REQUIRE_TRUE(c0->isSameShape(c0CorrectShape), 0,
|
|
"SRU operation: wrong shape of initial state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
|
|
if (mask)
|
|
REQUIRE_TRUE(mask->isSameShape(c0CorrectShape), 0,
|
|
"SRU operation: wrong shape of mask array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(mask).c_str());
|
|
|
|
// xm = x * mask
|
|
auto xm = x;
|
|
if (mask) {
|
|
xm = new NDArray(x->shapeInfo(), true, block.launchContext());
|
|
std::vector<LongType> dims = {0, 1};
|
|
x->applyBroadcast(broadcast::Multiply,&dims , mask, xm);
|
|
}
|
|
|
|
// time loop
|
|
helpers::sruTimeLoop(block.launchContext(), xm, c0, w, b, h, c);
|
|
|
|
if (mask) delete xm;
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
DECLARE_TYPES(sru) { getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS}); }
|
|
|
|
DECLARE_SHAPE_FN(sru) {
|
|
auto xShapeInfo = inputShape->at(0); // X, input 3d tensor [bS x inSize x time], time - number of time steps, bS -
|
|
// batch size, inSize - number of features
|
|
auto wShapeInfo = inputShape->at(1); // W, 2d tensor of weights [3*inSize x inSize]
|
|
auto bShapeInfo = inputShape->at(2); // B, row of biases with twice length [2*inSize]
|
|
auto c0ShapeInfo = inputShape->at(3); // C_{0}, 2d tensor of initial state [bS x inSize] at time t=0
|
|
auto maskShapeInfo =
|
|
block.width() > 4 ? inputShape->at(4) : nullptr; // optional, 2d tensor of dropout mask [bS x inSize]
|
|
|
|
const int rank = xShapeInfo[0]; // = 3
|
|
const int bS = xShapeInfo[1];
|
|
const int inSize = xShapeInfo[2];
|
|
const int time = xShapeInfo[3];
|
|
|
|
// input shapes validation
|
|
REQUIRE_TRUE(wShapeInfo[0] == rank - 1, 0,
|
|
"SRU operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
|
|
wShapeInfo[0]);
|
|
REQUIRE_TRUE(bShapeInfo[0] == 1, 0,
|
|
"SRU operation: wrong rank of biases array, expected is %i, but got %i instead !", 1, bShapeInfo[0]);
|
|
REQUIRE_TRUE(c0ShapeInfo[0] == rank - 1, 0,
|
|
"SRU operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
|
|
c0ShapeInfo[0]);
|
|
if (maskShapeInfo)
|
|
REQUIRE_TRUE(maskShapeInfo[0] == rank - 1, 0,
|
|
"SRU operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
|
|
maskShapeInfo[0]);
|
|
|
|
const std::vector<LongType> wCorrectShape = {3 * inSize, inSize};
|
|
const std::vector<LongType> bCorrectShape = {2 * inSize};
|
|
const std::vector<LongType> c0CorrectShape = {bS, inSize};
|
|
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
|
|
"SRU operation: wrong shape of weights array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
|
|
"SRU operation: wrong shape of biases array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, c0CorrectShape), 0,
|
|
"SRU operation: wrong shape of initial state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
|
|
if (maskShapeInfo)
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(maskShapeInfo, c0CorrectShape), 0,
|
|
"SRU operation: wrong shape of mask array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(maskShapeInfo).c_str());
|
|
|
|
LongType * newShapeInfo1 = nullptr;
|
|
ALLOCATE(newShapeInfo1, block.getWorkspace(), shape::shapeInfoLength(rank), sd::LongType); // [bS x inSize x time]
|
|
|
|
newShapeInfo1[0] = rank;
|
|
newShapeInfo1[1] = bS;
|
|
newShapeInfo1[2] = inSize;
|
|
newShapeInfo1[3] = time;
|
|
|
|
ShapeUtils::updateStridesAndType(newShapeInfo1, xShapeInfo, shape::order(xShapeInfo));
|
|
auto result = ConstantShapeHelper::getInstance().bufferForShapeInfo(newShapeInfo1)->primary();
|
|
RELEASE(newShapeInfo1, block.getWorkspace());
|
|
return SHAPELIST(result, result);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
CUSTOM_OP_IMPL(sru_bp, 8, 4, true, 0, 0) {
|
|
auto x = INPUT_VARIABLE(0);
|
|
// X, input 3d tensor [bS x K x N], N - number of time steps, bS - batch size, K - number of features
|
|
auto w = INPUT_VARIABLE(1); // W, 2d tensor of weights [3K x K]
|
|
auto b = INPUT_VARIABLE(2); // B, row of biases with twice length [1 x 2*K]
|
|
auto c0 = INPUT_VARIABLE(3); // C_{0}, 2d tensor of initial state [bS x K] at time t=0
|
|
auto c = INPUT_VARIABLE(4); // C, [bS x K x N]
|
|
auto inGradCt = INPUT_VARIABLE(5); // [bS x K]
|
|
auto inGradH = INPUT_VARIABLE(6); // [bS x K x N]
|
|
NDArray* mask = nullptr; // optional, 2d tensor of dropout mask [bS x K]
|
|
|
|
bool applyMask = false;
|
|
if (block.width() > 7) {
|
|
mask = INPUT_VARIABLE(7);
|
|
applyMask = true;
|
|
}
|
|
|
|
auto gradX = OUTPUT_VARIABLE(0); // [bS x K x N]
|
|
auto gradW = OUTPUT_VARIABLE(1); // [bS x 3K x K]
|
|
auto gradB = OUTPUT_VARIABLE(2); // [1 x 2K]
|
|
auto gradInit = OUTPUT_VARIABLE(3); // [bS x K]
|
|
|
|
const int bS = x->shapeOf()[0];
|
|
const int K = x->shapeOf()[1];
|
|
const int N = x->shapeOf()[2]; // N - number of time steps
|
|
|
|
std::vector<sd::LongType> gradBiasShape = {bS, 2 * K, N};
|
|
std::vector<sd::LongType> gradUShape = {bS, 3 * K, N};
|
|
std::vector<sd::LongType> gradHXShape = {bS, K, N};
|
|
std::vector<sd::LongType> gctShape = {bS, K};
|
|
std::vector<sd::LongType> gradTanhShape = {bS, K};
|
|
std::vector<sd::LongType> gradCtShape = {bS, K};
|
|
std::vector<sd::LongType> ftMinusShape = {bS, K};
|
|
std::vector<sd::LongType> rtMinusShape = {bS, K};
|
|
std::vector<sd::LongType> temp1Shape = {bS, K};
|
|
std::vector<sd::LongType> temp2Shape = {bS, K};
|
|
|
|
auto gradBias = NDArrayFactory::create_(x->ordering(), gradBiasShape, gradX->dataType(), block.launchContext());
|
|
auto gradU = NDArrayFactory::create_(x->ordering(), gradUShape, gradX->dataType(), block.launchContext());
|
|
auto gradHX = NDArrayFactory::create_(x->ordering(), gradHXShape, gradX->dataType(), block.launchContext());
|
|
auto gct = NDArrayFactory::create_(c->ordering(), gctShape, gradX->dataType(), block.launchContext());
|
|
auto gradTanh = NDArrayFactory::create_(c->ordering(), gradTanhShape, gradX->dataType(), block.launchContext());
|
|
auto gradCt = NDArrayFactory::create_(c->ordering(), gradCtShape, gradX->dataType(), block.launchContext());
|
|
auto ftMinus = NDArrayFactory::create_(c->ordering(), ftMinusShape, gradX->dataType(), block.launchContext());
|
|
auto rtMinus = NDArrayFactory::create_(c->ordering(), rtMinusShape, gradX->dataType(), block.launchContext());
|
|
auto temp1 = NDArrayFactory::create_(c->ordering(), temp1Shape, gradX->dataType(), block.launchContext());
|
|
auto temp2 = NDArrayFactory::create_(c->ordering(), temp2Shape, gradX->dataType(), block.launchContext());
|
|
std::vector<LongType> axes = {0, 1};
|
|
// x = x * mask
|
|
if (applyMask) x->applyBroadcast(broadcast::Multiply, &axes, mask, x); // apply mask
|
|
// multiplication matrix wi = matmul(w,x), U = WX
|
|
auto wi = MmulHelper::mmul(w, x, nullptr, 1., 0.); // U [bS x 3K x N]
|
|
|
|
auto wiZ = (*wi)({0, 0, 0, K, 0, 0}, true); // [bS x K x N]
|
|
auto wiF = (*wi)({0, 0, K, 2 * K, 0, 0}, true); // forget gate [bS x K x N]
|
|
auto wiR = (*wi)({0, 0, 2 * K, 3 * K, 0, 0}, true); // reset gate [bS x K x N]
|
|
auto bF = (*b)({0, 0, 0, K}, true); // biases for forget gate [1 x K]
|
|
auto bR = (*b)({0, 0, K, 2 * K}, true); // biases for reset gate [1 x K]
|
|
auto gradBF = (*gradBias)({0, 0, 0, K, 0, 0}, true); // [bS x K x N]
|
|
auto gradBR = (*gradBias)({0, 0, K, 2 * K, 0, 0}, true); // [bS x K x N]
|
|
auto gradUZ = (*gradU)({0, 0, 0, K, 0, 0}, true); // [bS x K x N]
|
|
auto gradUF = (*gradU)({0, 0, K, 2 * K, 0, 0}, true); // [bS x K x N]
|
|
auto gradUR = (*gradU)({0, 0, 2 * K, 3 * K, 0, 0}, true); // [bS x K x N]
|
|
|
|
NDArray* ct_1 = nullptr;
|
|
|
|
std::vector<LongType> idx = {0, 0, 0, 0, 0, 0};
|
|
|
|
for (int t = N - 1; t >= 0; --t) {
|
|
// initialization
|
|
idx[4] = t;
|
|
idx[5] = t + 1;
|
|
auto xt = (*x)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto zt = (*wiZ)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto ft = (*wiF)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto rt = (*wiR)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto ct = (*c)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto inGradHt = (*inGradH)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto gradBRt = (*gradBR)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto gradBFt = (*gradBF)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto gradHXt = (*gradHX)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto gradUZt = (*gradUZ)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto gradUFt = (*gradUF)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
auto gradURt = (*gradUR)(idx); // [bS x K x N] -> [bS x K x 1] -> [bS x K]
|
|
|
|
|
|
|
|
if (t != 0) {
|
|
idx[4] = t - 1;
|
|
idx[5] = t;
|
|
ct_1 = new NDArray((*c)(idx)); // previous c_{t-1}
|
|
} else
|
|
ct_1 = c0;
|
|
|
|
///////////////// forward
|
|
// ft = sigmoid(ft + bf), rt = sigmoid(rt + bR)
|
|
ft->addRowVector(bF, ft);
|
|
rt->addRowVector(bR, rt);
|
|
ft->applyTransform(transform::Sigmoid, ft);
|
|
rt->applyTransform(transform::Sigmoid, rt);
|
|
|
|
// TODO T val = (activation_type == 1) ? tanh(cur) : ((activation_type == 2) ? reluf(cur) : cur );
|
|
ct->applyTransform(transform::Tanh, gct);
|
|
// ftMinus = 1-ft, rtMinus = 1-rt
|
|
ft->applyTransform(transform::OneMinus, ftMinus);
|
|
rt->applyTransform(transform::OneMinus, rtMinus);
|
|
|
|
///////////////// backward
|
|
// bR, *grad_brt_ptr = inGradHt * (g_ct - xt) * (1.0f - rt) * rt;
|
|
gct->applyPairwiseTransform(pairwise::Subtract, xt, temp1); // temp1 = (g_ct - xt)
|
|
rtMinus->applyPairwiseTransform(pairwise::Multiply, rt, temp2); // temp2 = (1.0f - rt) * rt;
|
|
temp1->applyPairwiseTransform(pairwise::Multiply, temp2); // temp1 = (g_ct - xt) * (1.0f - rt) * rt;
|
|
inGradHt->applyPairwiseTransform(pairwise::Multiply, temp1,
|
|
gradBRt); // = inGradHt * (g_ct - xt) * (1.0f - rt) * rt;
|
|
|
|
// bF, TODO - tanh
|
|
// gradTanh = (1.0f - g_ct * g_ct);
|
|
gct->applyPairwiseTransform(pairwise::Multiply, gct, gradTanh); // gradTanh = g_ct * g_ct
|
|
gradTanh->applyTransform(transform::OneMinus, gradTanh); // gradTanh = (1.0f - g_ct * g_ct)
|
|
// gradCt = inGradHt * rt * gradTanh
|
|
rt->applyPairwiseTransform(pairwise::Multiply, gradTanh, gradCt); // gradCt = rt * gradTanh
|
|
inGradHt->applyPairwiseTransform(pairwise::Multiply, gradCt, gradCt); // gradCt = inGradHt * rt * gradTanh
|
|
// gradBFt = (gradCt + inGradCt) * (ct_1 - zt) * (1 - ft) * ft;
|
|
gradCt->applyPairwiseTransform(pairwise::Add, inGradCt, temp1); // temp1 = (gradCt + inGradCt)
|
|
ct_1->applyPairwiseTransform(pairwise::Subtract, zt, temp2); // temp2 = (ct_1 - zt)
|
|
temp1->applyPairwiseTransform(pairwise::Multiply, ftMinus, temp1); // temp1 = (gradCt + inGradCt)*(1-ft)
|
|
temp1->applyPairwiseTransform(pairwise::Multiply, ft, temp1); // temp1 = (gradCt + inGradCt)*(1-ft)*ft
|
|
temp1->applyPairwiseTransform(pairwise::Multiply, temp2,
|
|
gradBFt); // gradBFt = (gradCt + inGradCt) * (ct_1 - zt) * (1 - ft) * ft;
|
|
|
|
// x_t (highway connection), gradHXt = inGradHt * (1.0f - rt);
|
|
inGradHt->applyPairwiseTransform(pairwise::Multiply, rtMinus, gradHXt);
|
|
|
|
// U_t, gradUZt = (inGradHt * rt * grad_tanh + inGradCt) * (1.0f - ft);
|
|
rt->applyPairwiseTransform(pairwise::Multiply, gradTanh, temp1); // temp1 = rt * grad_tanh
|
|
inGradHt->applyPairwiseTransform(pairwise::Multiply, temp1,temp1); // temp1 = inGradHt * rt * grad_tanh
|
|
temp1->applyPairwiseTransform(pairwise::Add, inGradCt, temp1); // temp1 = inGradHt * rt * grad_tanh + inGradCt
|
|
temp1->applyPairwiseTransform(pairwise::Multiply, ftMinus,
|
|
gradUZt); // gradUZt = (inGradHt * rt * grad_tanh + inGradCt) * (1.0f - ft);
|
|
gradUFt->assign(gradBFt);
|
|
gradURt->assign(gradBRt);
|
|
|
|
// c_{t-1}, inGradCt = (gradCt + inGradCt) * ft;
|
|
gradCt->applyPairwiseTransform(pairwise::Add, inGradCt, temp1); // temp1 = (gradCt + inGradCt)
|
|
temp1->applyPairwiseTransform(pairwise::Multiply, ft, inGradCt); // inGradCt = (gradCt + inGradCt) * ft;
|
|
|
|
if (t != 0) delete ct_1;
|
|
|
|
delete xt;
|
|
delete zt;
|
|
delete ft;
|
|
delete rt;
|
|
delete ct;
|
|
delete inGradHt;
|
|
delete gradBRt;
|
|
delete gradHXt;
|
|
delete gradUZt;
|
|
delete gradUFt;
|
|
delete gradURt;
|
|
}
|
|
|
|
// gradInit
|
|
gradInit->assign(inGradCt);
|
|
|
|
// gradX
|
|
auto weightsT = w->transpose(); // [K x 3K]
|
|
MmulHelper::mmul(weightsT, gradU, gradX, 1., 0.); // [bS x K x N]
|
|
gradX->applyPairwiseTransform(pairwise::Add, gradHX, gradX);
|
|
std::vector<LongType> axes3 = {0, 1};
|
|
// + grad_highway_x
|
|
if (applyMask) gradX->applyBroadcast(broadcast::Multiply, &axes3, mask, gradX); // apply mask
|
|
|
|
// gradB
|
|
std::vector<sd::LongType> gradBShape = { 2 * K};
|
|
auto gradB2 = gradB->reshape(gradB->ordering(), gradBShape);
|
|
std::vector<LongType> axes2;
|
|
axes.push_back(0);
|
|
axes.push_back(2);
|
|
gradBias->reduceAlongDimension(reduce::Sum, gradB2, &axes2); // [1 x 2K]
|
|
|
|
// gradW [bS x 3K x K]
|
|
x->permutei({0, 2, 1}, false, false); // [bS x N x K]
|
|
MmulHelper::mmul(gradU, x, gradW, 1., 0.); // [bS x 3K x K]
|
|
|
|
delete gct;
|
|
delete gradU;
|
|
delete gradHX;
|
|
delete temp1;
|
|
delete temp2;
|
|
delete gradCt;
|
|
delete wi;
|
|
delete gradTanh;
|
|
delete ftMinus;
|
|
delete rtMinus;
|
|
delete gradBias;
|
|
delete weightsT;
|
|
return Status::OK;
|
|
}
|
|
|
|
DECLARE_TYPES(sru_bp) {
|
|
getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
|
|
}
|
|
|
|
DECLARE_SHAPE_FN(sru_bp) {
|
|
auto inShape = inputShape->at(0); // [bS x inSize x time]
|
|
auto bS = inShape[1];
|
|
auto inSize = inShape[2];
|
|
auto time = inShape[3];
|
|
char order = (char)(inShape[9]);
|
|
|
|
auto ret = SHAPELIST(
|
|
ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
|
|
std::vector<sd::LongType>{bS, inSize, time})->primary(),
|
|
ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
|
|
std::vector<sd::LongType>{bS, 3 * inSize, inSize})->primary(),
|
|
ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
|
|
std::vector<sd::LongType>{1, 2 * inSize})->primary(),
|
|
ConstantShapeHelper::getInstance().bufferForShapeInfo(ArrayOptions::dataType(inShape), order,
|
|
std::vector<sd::LongType>{bS, inSize})->primary()
|
|
);
|
|
return ret;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
CUSTOM_OP_IMPL(sru_bi, 5, 2, true, 0, 0) {
|
|
auto x = INPUT_VARIABLE(0); // X, input 3d tensor [time x bS x 2*inSize], time - number of time steps, bS - batch
|
|
// size, inSize - number of features
|
|
auto w = INPUT_VARIABLE(1); // W, 2d tensor of weights [2*inSize x 6*inSize]
|
|
auto b = INPUT_VARIABLE(2); // B, row of biases with twice length [1 x 4*inSize]
|
|
auto c0 = INPUT_VARIABLE(3); // C_{0}, 2d tensor of initial state [bS x 2*inSize] at time t=0
|
|
NDArray* mask =
|
|
block.width() > 4 ? INPUT_VARIABLE(4) : nullptr; // optional, 2d tensor of dropout mask [bS x 2*inSize]
|
|
|
|
auto ht = OUTPUT_VARIABLE(0); // h_t, [time x bS x 2*inSize]
|
|
auto ct = OUTPUT_VARIABLE(1); // c_t, [time x bS x 2*inSize]
|
|
|
|
// input shapes validation
|
|
const int rank = x->rankOf();
|
|
const LongType bS = x->sizeAt(1);
|
|
const LongType inSize = x->sizeAt(2) / 2;
|
|
|
|
REQUIRE_TRUE(x->rankOf() == rank, 0,
|
|
"SRU_BI operation: wrong rank of input array, expected is %i, but got %i instead !", rank, x->rankOf());
|
|
REQUIRE_TRUE(w->rankOf() == rank - 1, 0,
|
|
"SRU_BI operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
|
|
w->rankOf());
|
|
REQUIRE_TRUE(b->rankOf() == 1, 0, "SRU_BI operation: wrong rank of biases array, expected is 1, but got %i instead !",
|
|
b->rankOf());
|
|
REQUIRE_TRUE(c0->rankOf() == rank - 1, 0,
|
|
"SRU_BI operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
|
|
c0->rankOf());
|
|
if (mask)
|
|
REQUIRE_TRUE(mask->rankOf() == rank - 1, 0,
|
|
"SRU_BI operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
|
|
mask->rankOf());
|
|
|
|
const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
|
|
const std::vector<LongType> bCorrectShape = {4 * inSize};
|
|
const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
|
|
|
|
REQUIRE_TRUE(w->isSameShape(wCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
|
|
REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of biases array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
|
|
REQUIRE_TRUE(c0->isSameShape(c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
|
|
if (mask)
|
|
REQUIRE_TRUE(mask->isSameShape(c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(mask).c_str());
|
|
|
|
helpers::sruBI(block.launchContext(), x, w, b, c0, mask, ht, ct);
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
DECLARE_TYPES(sru_bi) {
|
|
getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
|
|
}
|
|
|
|
DECLARE_SHAPE_FN(sru_bi) {
|
|
auto xShapeInfo = inputShape->at(0); // [time x bS x 2K ]
|
|
auto wShapeInfo = inputShape->at(1);
|
|
auto bShapeInfo = inputShape->at(2);
|
|
auto c0ShapeInfo = inputShape->at(3);
|
|
auto maskShapeInfo =
|
|
block.width() > 4 ? inputShape->at(4) : nullptr; // optional, 2d tensor of dropout mask [bS x inSize]
|
|
|
|
const int rank = xShapeInfo[0]; // = 3
|
|
const LongType time = xShapeInfo[1];
|
|
const LongType bS = xShapeInfo[2];
|
|
const LongType inSize = xShapeInfo[3] / 2;
|
|
|
|
// input shapes validation
|
|
REQUIRE_TRUE(wShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
|
|
wShapeInfo[0]);
|
|
REQUIRE_TRUE(bShapeInfo[0] == 1, 0,
|
|
"SRU_BI operation: wrong rank of biases array, expected is 1, but got %i instead !", bShapeInfo[0]);
|
|
REQUIRE_TRUE(c0ShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
|
|
c0ShapeInfo[0]);
|
|
if (maskShapeInfo)
|
|
REQUIRE_TRUE(maskShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
|
|
maskShapeInfo[0]);
|
|
|
|
const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
|
|
const std::vector<LongType> bCorrectShape = {4 * inSize};
|
|
const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
|
|
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of biases array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
|
|
if (maskShapeInfo)
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(maskShapeInfo, c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(maskShapeInfo).c_str());
|
|
|
|
char order = shape::order(xShapeInfo);
|
|
|
|
ShapeDescriptor *descriptor = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {time, bS, 2 * inSize});
|
|
auto result = ConstantShapeHelper::getInstance().createShapeInfo(descriptor);
|
|
return SHAPELIST(result, result);
|
|
}
|
|
|
|
DECLARE_TYPES(sru_bi_bp) {
|
|
getOpDescriptor()->setAllowedInputTypes(ANY)->setAllowedOutputTypes({ALL_FLOATS});
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
CUSTOM_OP_IMPL(sru_bi_bp, 8, 4, true, 0, 0) {
|
|
auto x = INPUT_VARIABLE(0); // X, input 3d tensor [time x bS x 2*inSize], time - number of time steps, bS - batch
|
|
// size, inSize - number of features
|
|
auto w = INPUT_VARIABLE(1); // W, 2d tensor of weights [2*inSize x 6*inSize]
|
|
auto b = INPUT_VARIABLE(2); // B, row of biases with twice length [4*inSize]
|
|
auto c0 = INPUT_VARIABLE(3); // C_{0}, 2d tensor of initial state [bS x 2*inSize] at time t=0
|
|
auto ct = INPUT_VARIABLE(4); // C, [time x bS x 2*inSize]
|
|
auto inGradC0 = INPUT_VARIABLE(5); // [bS x 2*inSize]
|
|
auto inGradHt = INPUT_VARIABLE(6); // [time x bS x 2*inSize]
|
|
NDArray* mask =
|
|
block.width() > 7 ? INPUT_VARIABLE(7) : nullptr; // optional, 2d tensor of dropout mask [bS x 2*inSize]
|
|
|
|
// input shapes validation
|
|
const int rank = x->rankOf();
|
|
const LongType time = x->sizeAt(0);
|
|
const LongType bS = x->sizeAt(1);
|
|
const LongType inSize = x->sizeAt(2) / 2;
|
|
|
|
REQUIRE_TRUE(w->rankOf() == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
|
|
w->rankOf());
|
|
REQUIRE_TRUE(b->rankOf() == 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of biases array, expected is 1, but got %i instead !", b->rankOf());
|
|
REQUIRE_TRUE(c0->rankOf() == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
|
|
c0->rankOf());
|
|
REQUIRE_TRUE(ct->rankOf() == rank, 0,
|
|
"SRU_BI_BP operation: wrong rank of state array, expected is %i, but got %i instead !", rank,
|
|
ct->rankOf());
|
|
REQUIRE_TRUE(inGradC0->rankOf() == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of gradient c0, expected is %i, but got %i instead !", rank - 1,
|
|
inGradC0->rankOf());
|
|
REQUIRE_TRUE(inGradHt->rankOf() == rank, 0,
|
|
"SRU_BI_BP operation: wrong rank of gradient ht, expected is %i, but got %i instead !", rank,
|
|
inGradHt->rankOf());
|
|
if (mask)
|
|
REQUIRE_TRUE(mask->rankOf() == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
|
|
mask->rankOf());
|
|
|
|
const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
|
|
const std::vector<LongType> bCorrectShape = {4 * inSize};
|
|
const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
|
|
const std::vector<LongType> ctCorrectShape = {time, bS, 2 * inSize};
|
|
|
|
REQUIRE_TRUE(w->isSameShape(wCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(w).c_str());
|
|
REQUIRE_TRUE(b->isSameShape(bCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of biases array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(b).c_str());
|
|
REQUIRE_TRUE(c0->isSameShape(c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0).c_str());
|
|
REQUIRE_TRUE(ct->isSameShape(ctCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(ctCorrectShape).c_str(), ShapeUtils::shapeAsString(ct).c_str());
|
|
if (mask)
|
|
REQUIRE_TRUE(mask->isSameShape(c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(mask).c_str());
|
|
|
|
auto gradI = OUTPUT_VARIABLE(0); // [time x bS x 2*inSize]
|
|
auto gradW = OUTPUT_VARIABLE(1); // [time x 2*inSize x 6*inSize]
|
|
auto gradB = OUTPUT_VARIABLE(2); // [1 x 4*inSize]
|
|
auto gradC0 = OUTPUT_VARIABLE(3); // [bS x 2*inSize]
|
|
|
|
helpers::sruBIBP(block.launchContext(), x, w, b, c0, ct, inGradC0, inGradHt, mask, gradI, gradW, gradB, gradC0);
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
DECLARE_SHAPE_FN(sru_bi_bp) {
|
|
auto xShapeInfo = inputShape->at(0); // [time x bS x 2K ]
|
|
auto wShapeInfo = inputShape->at(1);
|
|
auto bShapeInfo = inputShape->at(2);
|
|
auto c0ShapeInfo = inputShape->at(3);
|
|
auto ctShapeInfo = inputShape->at(4);
|
|
auto inGradC0ShapeInfo = inputShape->at(5);
|
|
auto inGradHtShapeInfo = inputShape->at(6);
|
|
auto maskShapeInfo =
|
|
block.width() > 7 ? inputShape->at(7) : nullptr; // optional, 2d tensor of dropout mask [bS x inSize]
|
|
|
|
// input shapes validation
|
|
const int rank = xShapeInfo[0];
|
|
const LongType time = xShapeInfo[1];
|
|
const LongType bS = xShapeInfo[2];
|
|
const LongType inSize = xShapeInfo[3] / 2;
|
|
|
|
REQUIRE_TRUE(wShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of weights array, expected is %i, but got %i instead !", rank - 1,
|
|
wShapeInfo[0]);
|
|
REQUIRE_TRUE(bShapeInfo[0] == 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of biases array, expected is 1, but got %i instead !", bShapeInfo[0]);
|
|
REQUIRE_TRUE(c0ShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of initial state array, expected is %i, but got %i instead !", rank - 1,
|
|
c0ShapeInfo);
|
|
REQUIRE_TRUE(ctShapeInfo[0] == rank, 0,
|
|
"SRU_BI_BP operation: wrong rank of state array, expected is %i, but got %i instead !", rank,
|
|
ctShapeInfo);
|
|
REQUIRE_TRUE(inGradC0ShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of gradient c0, expected is %i, but got %i instead !", rank - 1,
|
|
inGradC0ShapeInfo[0]);
|
|
REQUIRE_TRUE(inGradHtShapeInfo[0] == rank, 0,
|
|
"SRU_BI_BP operation: wrong rank of gradient ht, expected is %i, but got %i instead !", rank,
|
|
inGradHtShapeInfo[0]);
|
|
if (maskShapeInfo)
|
|
REQUIRE_TRUE(maskShapeInfo[0] == rank - 1, 0,
|
|
"SRU_BI_BP operation: wrong rank of mask array, expected is %i, but got %i instead !", rank - 1,
|
|
maskShapeInfo[0]);
|
|
|
|
const std::vector<LongType> wCorrectShape = {2 * inSize, 6 * inSize};
|
|
const std::vector<LongType> bCorrectShape = {4 * inSize};
|
|
const std::vector<LongType> c0CorrectShape = {bS, 2 * inSize};
|
|
const std::vector<LongType> ctCorrectShape = {time, bS, 2 * inSize};
|
|
const std::vector<LongType> inGradC0CorrectShape = {bS, 2 * inSize};
|
|
const std::vector<LongType> inGradHtCorrectShape = {time, bS, 2 * inSize};
|
|
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(wShapeInfo, wCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of weights array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(wCorrectShape).c_str(), ShapeUtils::shapeAsString(wShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(bShapeInfo, bCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of biases array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(bCorrectShape).c_str(), ShapeUtils::shapeAsString(bShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(c0ShapeInfo, c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of initial state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(c0ShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(ctShapeInfo, ctCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of state array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(ctCorrectShape).c_str(), ShapeUtils::shapeAsString(ctShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(inGradC0ShapeInfo, inGradC0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of gradient c0 array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(inGradC0CorrectShape).c_str(),
|
|
ShapeUtils::shapeAsString(inGradC0ShapeInfo).c_str());
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(inGradHtShapeInfo, inGradHtCorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of gradient ht array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(inGradHtCorrectShape).c_str(),
|
|
ShapeUtils::shapeAsString(inGradHtShapeInfo).c_str());
|
|
if (maskShapeInfo)
|
|
REQUIRE_TRUE(ShapeUtils::areShapesEqual(maskShapeInfo, c0CorrectShape), 0,
|
|
"SRU_BI operation: wrong shape of mask array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(c0CorrectShape).c_str(), ShapeUtils::shapeAsString(maskShapeInfo).c_str());
|
|
|
|
const char order = shape::order(xShapeInfo);
|
|
|
|
ShapeDescriptor *descriptor1 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {time, bS, 2 * inSize});
|
|
ShapeDescriptor *descriptor2 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {time, 2 * inSize, 6 * inSize});
|
|
ShapeDescriptor *descriptor3 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {4 * inSize});
|
|
ShapeDescriptor *descriptor4 = new ShapeDescriptor(ArrayOptions::dataType(xShapeInfo), order, {bS, 2 * inSize});
|
|
|
|
return SHAPELIST(ConstantShapeHelper::getInstance().createShapeInfo(descriptor1),
|
|
ConstantShapeHelper::getInstance().createShapeInfo(descriptor2),
|
|
ConstantShapeHelper::getInstance().createShapeInfo(descriptor3),
|
|
ConstantShapeHelper::getInstance().createShapeInfo(descriptor4));
|
|
}
|
|
|
|
} // namespace ops
|
|
} // namespace sd
|
|
|
|
#endif
|