cavis/libnd4j/include/ops/declarable/helpers/impl/rnn.cpp

105 lines
4.0 KiB
C++
Raw Normal View History

2019-06-06 14:21:15 +02:00
/*******************************************************************************
* Copyright (c) 2015-2018 Skymind, Inc.
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com), created on 16.04.2018
//
// function nnCell implements an Elman RNN cell: output = activation(Wx*x + bx + Wh*ht + bh)
#include<ops/declarable/helpers/rnn.h>
#include <helpers/BlasHelper.h>
namespace nd4j {
namespace ops {
namespace helpers {
//////////////////////////////////////////////////////////////////////////
static FORCEINLINE NDArray activation(const NDArray& arr) {
return (const_cast<NDArray&>(arr)).transform(transform::Tanh);
}
//////////////////////////////////////////////////////////////////////////
void rnnCell(nd4j::LaunchContext * context, const NDArray* xt, const NDArray* Wx, const NDArray* Wh, const NDArray* b, const NDArray* ht_1, NDArray* ht) {
// xt input [bS x inSize]
// Wx input-to-hidden weights, [inSize x numUnits]
// Wh hidden-to-hidden weights, [numUnits x numUnits]
// b biases, [2*numUnits]: {0, numUnits} are input-to-hidden biases and {numUnits, 2*numUnits} are hidden-to-hidden biases
// ht_1 previous cell output [bS x numUnits], that is at previous time step t-1, in case of projection=false -> numUnits=numUnits!!!
const int numUnits = ht_1->sizeAt(1);
// ht is current cell output [bS x numUnits], that is at current time step t
ht->assign(activation(mmul(*xt, *Wx) + (*b)({{0, numUnits}}) + mmul(*ht_1, *Wh) + (*b)({{numUnits, 2*numUnits}}))); // [bS x numUnits] + [numUnits] + [bS x numUnits] + [numUnits] = [bS x numUnits]
}
//////////////////////////////////////////////////////////////////////////
void rnnTimeLoop(nd4j::LaunchContext * context, const NDArray* x, const NDArray* Wx, const NDArray* Wh, const NDArray* b, const NDArray* h0, const NDArray* maxTimeStep, NDArray* h, NDArray* hFinal) {
// x input [time x bS x inSize]
// Wx input-to-hidden weights, [inSize x numUnits]
// Wh hidden-to-hidden weights, [numUnits x numUnits]
// b biases for, [2*numUnits]
// h0 initial cell output (at time step = 0) [bS x numUnits]
// maxTimeStep vector [bS] containing integer values within [0,time), each element of this vector set max time step per each input in batch, this means there are no calculations for time >= maxTimeStep
const int time = x->sizeAt(0);
const int bS = x->sizeAt(1);
// at first time step
if(h0)
hFinal->assign(h0);
else
*hFinal = 0.;
BlasHelper::getInstance(); // to avoid memory leak in pragma parallel loops
// loop through batch of inputs
for (int e = 0; e < bS; ++e) {
// loop through time steps
for (int t = 0; t < time; ++t) {
int maxStep = maxTimeStep ? maxTimeStep->e<int>(e) : time;
auto xt = (*x)({t,t+1, e,e+1, 0,0}, true);
auto ht = (*h)({t,t+1, e,e+1, 0,0}, true);
auto ht_1 = (*hFinal)({e,e+1, 0,0}, true); // previous state
if(t >= maxStep) {
ht = 0.;
if(maxStep != 0)
ht_1.assign((*h)({maxStep-1,maxStep, e,e+1, 0,0}));
}
else {
helpers::rnnCell(context, &xt, Wx, Wh, b, &ht_1, &ht);
ht_1.assign(ht);
}
}
}
}
}
}
}