From 4746121d438c72d0287a6ff7edb193a2da68029d Mon Sep 17 00:00:00 2001
From: Francesco Gatti <gattifrancesco@hotmail.it>
Date: Sat, 15 Feb 2020 20:37:08 +0100
Subject: [PATCH] LSTM params

---
 include/tkDNN/Layer.h     |  5 ++-
 src/LSTM.cpp              | 94 +++++++++++++++++++++++++++++++++------
 tests/imuodom/imuodom.cpp | 20 ++++++---
 tests/imuodom/infer.py    |  1 +
 tests/weights_exporter.py | 16 +++++--
 5 files changed, 111 insertions(+), 25 deletions(-)

diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h
index 240d9a2..319df6a 100644
--- a/include/tkDNN/Layer.h
+++ b/include/tkDNN/Layer.h
@@ -212,6 +212,7 @@ protected:
     https://github.com/jiangnanhugo/seq2seq_cuda/blob/e4dbdcfa0517c972bfd4beea9f11a5233954093c/src/rnn.cpp
     https://github.com/Jeffery-Song/mxnet-test/blob/aab666faad44011f7a67b527b5f6c960367d0422/src/operator/cudnn_rnn-inl.h
     https://stackoverflow.com/a/38737941
+    https://colah.github.io/posts/2015-08-Understanding-LSTMs/
 
     PARAMS (numlayers*2):
         layer0:
@@ -236,7 +237,7 @@ public:
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 
-    const bool bidirectional = 1; /**> is the net bidir */
+    const bool bidirectional = false; /**> 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 */
@@ -254,9 +255,11 @@ protected:
     cudnnTensorDescriptor_t hx_desc_, cx_desc_;
     cudnnTensorDescriptor_t hy_desc_, cy_desc_;
     dnnType *hx_ptr, *cx_ptr, *hy_ptr, *cy_ptr;
+    int stateDataDim;
 
     cudnnFilterDescriptor_t w_desc_;
     dnnType *w_ptr;
+    dnnType *w_h;
 };
 
 
diff --git a/src/LSTM.cpp b/src/LSTM.cpp
index c31bccf..ce2a9f5 100644
--- a/src/LSTM.cpp
+++ b/src/LSTM.cpp
@@ -66,10 +66,12 @@ LSTM::LSTM( Network *net, int hiddensize, bool returnSeq, std::string fname_weig
     checkCUDNN(cudnnSetTensorNdDescriptor(hy_desc_, net->dataType, 3, dimA, strideA));
     checkCUDNN(cudnnSetTensorNdDescriptor(cy_desc_, net->dataType, 3, dimA, strideA));
     // allocate     dnnType *hx_ptr, *cx_ptr, *hy_ptr, *cy_ptr;
-    checkCuda( cudaMalloc(&hx_ptr, dimA[0]*dimA[1]*dimA[2]*sizeof(dnnType)) );
-    checkCuda( cudaMalloc(&cx_ptr, dimA[0]*dimA[1]*dimA[2]*sizeof(dnnType)) );
-    checkCuda( cudaMalloc(&hy_ptr, dimA[0]*dimA[1]*dimA[2]*sizeof(dnnType)) );
-    checkCuda( cudaMalloc(&cy_ptr, dimA[0]*dimA[1]*dimA[2]*sizeof(dnnType)) );
+    stateDataDim = dimA[0]*dimA[1]*dimA[2];
+    checkCuda( cudaMalloc(&hx_ptr, stateDataDim*sizeof(dnnType)) );
+    checkCuda( cudaMalloc(&cx_ptr, stateDataDim*sizeof(dnnType)) );
+    checkCuda( cudaMalloc(&hy_ptr, stateDataDim*sizeof(dnnType)) );
+    checkCuda( cudaMalloc(&cy_ptr, stateDataDim*sizeof(dnnType)) );
+    
 
 
     // Create Dropout descriptors // TODO: ??? IS IT NECESSARY ???
@@ -89,7 +91,7 @@ LSTM::LSTM( Network *net, int hiddensize, bool returnSeq, std::string fname_weig
     checkCUDNN(cudnnSetRNNDescriptor(net->cudnnHandle, 
         rnnDesc, stateSize, numLayers, dropoutDesc,
         cudnnRNNInputMode_t::CUDNN_LINEAR_INPUT,
-        cudnnDirectionMode_t::CUDNN_BIDIRECTIONAL,
+        (bidirectional ? cudnnDirectionMode_t::CUDNN_BIDIRECTIONAL : cudnnDirectionMode_t::CUDNN_UNIDIRECTIONAL),
         cudnnRNNMode_t::CUDNN_LSTM,
         cudnnRNNAlgo_t::CUDNN_RNN_ALGO_STANDARD,
         net->dataType));
@@ -115,22 +117,81 @@ LSTM::LSTM( Network *net, int hiddensize, bool returnSeq, std::string fname_weig
     dim_w[0] = cudnn_params;
     checkCUDNN(cudnnSetFilterNdDescriptor(w_desc_,
         net->dataType, net->tensorFormat, 3, dim_w));
-    // allocate params dnnType *w_ptr;
-    checkCuda( cudaMalloc(&w_ptr, cudnn_params*sizeof(dnnType)) );
-
 
+    // load params
+    readBinaryFile(fname_weights, cudnn_params, &w_h, &w_ptr);
     
     //allocate data for infer result
-    int dstDim = input_dim.n * stateSize*2 * input_dim.h * input_dim.w;
+    int dstDim = input_dim.n * stateSize*(bidirectional ? 2 : 1) * input_dim.h * input_dim.w;
     checkCuda( cudaMalloc(&dstData, dstDim*sizeof(dnnType)) );
 
     // set output dim 
     output_dim = input_dim;
-    output_dim.c = stateSize*2;
+    output_dim.c = stateSize*(bidirectional ? 2 : 1);
     if(!returnSeq) {
         output_dim.h = 1;
         output_dim.w = 1;
     }
+
+
+
+
+    // Query weight layout
+    cudnnFilterDescriptor_t m_desc;
+    checkCUDNN(cudnnCreateFilterDescriptor(&m_desc));
+    dnnType *p;
+    int n = 8; // lstm layers
+    
+    printCenteredTitle("WEIGHTS", '=', 20);
+    for (int i = 0; i < numLayers*(bidirectional?2:1); ++i) {
+        for (int j = 0; j < n; ++j) {
+            
+            checkCUDNN(cudnnGetRNNLinLayerMatrixParams(net->cudnnHandle, rnnDesc,
+                i, x_desc_vec_[0], w_desc_, 0, j, m_desc, (void**)&p));
+            
+            std::cout << "ptr: " << ((int64_t)(p - NULL))/sizeof(dnnType)<<"\n";
+            
+            cudnnDataType_t t;
+            cudnnTensorFormat_t f;
+            int ndim = 5;
+            int dims[5] = {0, 0, 0, 0, 0};
+            checkCUDNN(cudnnGetFilterNdDescriptor(m_desc, ndim, &t, &f, &ndim, &dims[0]));
+            std::cout << "(layer, linlayer): " <<  i << " " << j << "\n";
+
+            int tot = 1;
+            for (int i = 0; i < ndim; ++i) {
+                std::cout << dims[i] << " ";
+                tot *= dims[i];
+            }
+            std::cout<<"\t-> "<<tot<<"\n\n";
+        }
+    }
+
+    printCenteredTitle("BIAS", '=', 20);
+    for (int i = 0; i < numLayers*(bidirectional?2:1); ++i) {
+        for (int j = 0; j < n; ++j) {
+            checkCUDNN(cudnnGetRNNLinLayerBiasParams(net->cudnnHandle, rnnDesc, 
+                i, x_desc_vec_[0], w_desc_, 0, j, m_desc, (void**)&p));
+
+            std::cout << "ptr: " << ((int64_t)(p - NULL))/sizeof(dnnType)<<"\n";
+
+            cudnnDataType_t t;
+            cudnnTensorFormat_t f;
+            int ndim = 5;
+            int dims[5] = {0, 0, 0, 0, 0};
+            checkCUDNN(cudnnGetFilterNdDescriptor(m_desc, ndim, &t, &f, &ndim, &dims[0]));
+            std::cout << "(layer, linlayer): " <<  i << " " << j << "\n";
+
+            int tot = 1;
+            for (int i = 0; i < ndim; ++i) {
+                std::cout << dims[i] << " ";
+                tot *= dims[i];
+            }
+            std::cout<<"\t-> "<<tot<<"\n\n";
+        }
+    }
+    
+    checkCUDNN(cudnnDestroyFilterDescriptor(m_desc));
 }
 
 LSTM::~LSTM() {
@@ -149,18 +210,23 @@ LSTM::~LSTM() {
 dnnType* LSTM::infer(dataDim_t &dim, dnnType* srcData) {
     std::cout<<"LSTM infer\n";
 
+    // reset states
+    checkCuda( cudaMemset(hx_ptr, 0, stateDataDim*sizeof(float)) );	
+    checkCuda( cudaMemset(cx_ptr, 0, stateDataDim*sizeof(float)) );
+
+
     checkCUDNN(cudnnRNNForwardInference(net->cudnnHandle,
         rnnDesc,
-        seqLen,
-        x_desc_vec_.data(),         // input array of desc
+        seqLen,                     // number of time steps (nT)
+        x_desc_vec_.data(),         // input array of desc (nT*nC_in)
         srcData,                    // input pointer
-        hx_desc_,                   // initial hidden state desc    
+        hx_desc_,                   // initial hidden state desc     
         hx_ptr,                     // initial hidden state pointer 
         cx_desc_,                   // initial cell state desc      
         cx_ptr,                     // initial cell state pointer   
         w_desc_,                    // weights desc
         w_ptr,                      // weights pointer
-        y_desc_vec_.data(),         // output desc
+        y_desc_vec_.data(),         // output desc     (nT*nC_out)
         dstData,                    // output pointer
         hy_desc_,                   // final hidden state desc        
         hy_ptr,                     // final hidden state pointer 
diff --git a/tests/imuodom/imuodom.cpp b/tests/imuodom/imuodom.cpp
index 6fc3fba..4bb4d46 100644
--- a/tests/imuodom/imuodom.cpp
+++ b/tests/imuodom/imuodom.cpp
@@ -14,6 +14,9 @@ const char *c2_bin = "../tests/imuodom/layers/conv1d_9.bin";
 const char *c3_bin = "../tests/imuodom/layers/conv1d_10.bin";
 const char *c4_bin = "../tests/imuodom/layers/conv1d_11.bin";
 const char *c5_bin = "../tests/imuodom/layers/conv1d_12.bin";
+const char *l0_bin = "../tests/imuodom/layers/bidirectional_3.bin";
+const char *l1_bin = "../tests/imuodom/layers/bidirectional_4.bin";
+const char *d0_bin = "../tests/imuodom/layers/dense_3.bin";
 
 int main() {
 
@@ -48,8 +51,9 @@ 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, true, "ciao");
-    tk::dnn::LSTM lstm1(&net, 128, false, "ciao");
+    //tk::dnn::LSTM lstm0(&net, 128, true, l0_bin);
+    //tk::dnn::LSTM lstm1(&net, 128, false, l1_bin);
+    //tk::dnn::Dense d0 (&net, 3, d0_bin);
 
     net.print();
 
@@ -62,10 +66,12 @@ int main() {
     TIMER_STOP
 
     // Print real test
-    std::cout<<"\n==== CHECK RESULT ====\n";
-    dnnType *out;
-    dnnType *out_h;
-    readBinaryFile(output_bin, dim.tot(), &out_h, &out);
-    checkResult(dim.tot(), data, out);
+    //std::cout<<"\n==== CHECK RESULT ====\n";
+    //dnnType *out;
+    //dnnType *out_h;
+    //readBinaryFile(output_bin, dim.tot(), &out_h, &out);
+    //checkResult(dim.tot(), data, out);
+
+    printDeviceVector(100, data);
     return 0;
 }
diff --git a/tests/imuodom/infer.py b/tests/imuodom/infer.py
index 9d9929d..ec0e2a3 100644
--- a/tests/imuodom/infer.py
+++ b/tests/imuodom/infer.py
@@ -27,6 +27,7 @@ if __name__ == '__main__':
 
     weights = model.get_weights()
 
+    np.random.seed(2)
     x_angle = np.random.rand(1,100,4)
     x_gyro = np.random.rand(1,100,3)
     x_acc = np.random.rand(1,100,3)
diff --git a/tests/weights_exporter.py b/tests/weights_exporter.py
index 912ef82..fa07fe4 100644
--- a/tests/weights_exporter.py
+++ b/tests/weights_exporter.py
@@ -24,6 +24,8 @@ def export_layer(name, weights, bias):
         weights = weights.transpose(3,2,0,1)
     elif(weights.ndim == 3):
         weights = weights.transpose(2,1,0)
+    elif(weights.ndim == 2):
+        weights = weights.transpose(1,0)
     else:
         print("Ndim", weights.ndim)
         raise("not implemented with dim" )
@@ -40,13 +42,19 @@ def export_layer(name, weights, bias):
     bin_write(f, bias)
     print ("WEIGHTS saved\n")
 
-def export_bidir(name, weights):
+def export_bidir(name, params):
     print ("########    EXPORT", name, "LAYER    ########")
  
-    for w in weights:
-        print(np.shape(w))
+    f = open(name + ".bin", mode='wb')
 
 
+    for w in params:
+        #w = w.transpose()
+        print(np.shape(w))
+        bin_write(f, w)
+   
+    print("WEIGHTS saved\n")
+
 #https://github.com/fchollet/keras/wiki/Converting-convolution-kernels-from-Theano-to-TensorFlow-and-vice-versa
 if __name__ == '__main__':
     print("DATA FORMAT: ", keras.backend.image_data_format())
@@ -64,6 +72,7 @@ if __name__ == '__main__':
     model = load_model(args.model)
     model.summary()
 
+
     weights = model.get_weights()
 
     ws = np.shape(weights)
@@ -90,6 +99,7 @@ if __name__ == '__main__':
         elif name.startswith("dense"):
             export_layer(args.output + "/" + name, wgs[0], wgs[1])
         elif name.startswith("bidirectional"):
+            wgs = l.forward_layer.get_weights()
             export_bidir(args.output + "/" + name, wgs)
         else:
             print ("skip:", name, "has no weights")