diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h index 596fbec..240d9a2 100644 --- a/include/tkDNN/Layer.h +++ b/include/tkDNN/Layer.h @@ -230,13 +230,14 @@ protected: class LSTM : public Layer { public: - LSTM(Network *net, int hiddensize, std::string fname_weights); + LSTM(Network *net, int hiddensize, bool returnSeq, std::string fname_weights); virtual ~LSTM(); virtual layerType_t getLayerType() { return LAYER_LSTM; }; virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); const bool bidirectional = 1; /**> is the net bidir */ + bool returnSeq = false; /**> if false return only the result of last timestep */ int stateSize = 0; /**> number of hidden states */ int seqLen = 0; /**> number of timesteps */ int numLayers = 1; /**> number of internal layers */ diff --git a/src/LSTM.cpp b/src/LSTM.cpp index 46a0593..c31bccf 100644 --- a/src/LSTM.cpp +++ b/src/LSTM.cpp @@ -4,9 +4,10 @@ namespace tk { namespace dnn { -LSTM::LSTM( Network *net, int hiddensize, std::string fname_weights) : +LSTM::LSTM( Network *net, int hiddensize, bool returnSeq, std::string fname_weights) : Layer(net) { + this->returnSeq = returnSeq; int batchSize = input_dim.n; int inputSize = input_dim.c; seqLen = input_dim.w; @@ -117,12 +118,19 @@ LSTM::LSTM( Network *net, int hiddensize, std::string fname_weights) : // allocate params dnnType *w_ptr; checkCuda( cudaMalloc(&w_ptr, cudnn_params*sizeof(dnnType)) ); + + + //allocate data for infer result + int dstDim = input_dim.n * stateSize*2 * input_dim.h * input_dim.w; + checkCuda( cudaMalloc(&dstData, dstDim*sizeof(dnnType)) ); + // set output dim output_dim = input_dim; output_dim.c = stateSize*2; - - //allocate data for infer result - checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) ); + if(!returnSeq) { + output_dim.h = 1; + output_dim.w = 1; + } } LSTM::~LSTM() { diff --git a/tests/imuodom/imuodom.cpp b/tests/imuodom/imuodom.cpp index 14146b2..6fc3fba 100644 --- a/tests/imuodom/imuodom.cpp +++ b/tests/imuodom/imuodom.cpp @@ -48,7 +48,8 @@ int main() { tk::dnn::Layer *concat_l[3] = { &x0_2, &x1_2, &x2_2 }; tk::dnn::Route concat (&net, concat_l, 3); - tk::dnn::LSTM lstm0(&net, 128, "ciao"); + tk::dnn::LSTM lstm0(&net, 128, true, "ciao"); + tk::dnn::LSTM lstm1(&net, 128, false, "ciao"); net.print();