lstm return seq

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 */

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() {

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();