LSTM cudnn test

2020-02-13 19:27:18 +01:00
parent 6a3b261bc9
commit c876fa05ae
10 changed files with 464 additions and 134 deletions
@@ -0,0 +1,142 @@
+#include <iostream>
+
+#include "Layer.h"
+
+namespace tk { namespace dnn {
+
+LSTM::LSTM( Network *net, int hiddensize, std::string fname_weights) :
+    Layer(net) {
+
+    checkCUDNN( cudnnCreateFilterDescriptor(&paramDesc));
+    checkCUDNN( cudnnCreateRNNDescriptor(&rnnDesc) );
+    checkCUDNN( cudnnCreateRNNDataDescriptor(&rnnDataDesc) );
+    checkCUDNN( cudnnCreateDropoutDescriptor(&dropDesc));
+
+    int n = input_dim.n;
+    int c = input_dim.c;
+    int h = input_dim.h;
+    int w = input_dim.w;
+    checkCUDNN( cudnnSetTensor4dDescriptor(srcTensorDesc,
+                                           net->tensorFormat, net->dataType, n, 1, h, w) );
+
+    int numlayers = 1;
+    checkCUDNN( cudnnSetRNNDescriptor(net->cudnnHandle, rnnDesc, hiddensize, numlayers, dropDesc,
+            cudnnRNNInputMode_t::CUDNN_LINEAR_INPUT,
+            cudnnDirectionMode_t::CUDNN_BIDIRECTIONAL, cudnnRNNMode_t::CUDNN_LSTM,
+            cudnnRNNAlgo_t::CUDNN_RNN_ALGO_STANDARD, net->dataType) );
+
+    // find dimension of params
+    size_t params_size = 0;
+    checkCUDNN( cudnnGetRNNParamsSize(net->cudnnHandle, rnnDesc, srcTensorDesc, &params_size, net->dataType) );
+    std::cout<<"Params size bytes: "<<params_size<<", floats:  "<<params_size/4<<"\n";
+
+
+    int dimW[3] = { int(params_size / sizeof(float)), 1, 1};
+    checkCUDNN(cudnnCreateFilterDescriptor(&paramDesc));
+    checkCUDNN(cudnnSetFilterNdDescriptor(paramDesc, net->dataType, net->tensorFormat, 3, dimW));
+    checkCuda( cudaMalloc(&paramsSpace, params_size) );
+
+
+    int numlinearlayers = 8;
+
+    for(int i=0; i<numlayers*2; i++) {
+        std::cout<<"layer: "<<i<<"\n";
+        for(int j=0; j<numlinearlayers; j++) {
+
+            // get weights pointer
+            cudnnFilterDescriptor_t linLayerMatDesc;
+            checkCUDNN(cudnnCreateFilterDescriptor(&linLayerMatDesc));
+            dnnType *linLayerMat;
+
+            checkCUDNN(cudnnGetRNNLinLayerMatrixParams(net->cudnnHandle, rnnDesc,
+                    i, srcTensorDesc, paramDesc, paramsSpace,
+                    j, linLayerMatDesc, (void **)&linLayerMat));
+
+            if(linLayerMat == nullptr) {
+                FatalError("LSTM No weights in hidden layer");
+            }
+
+            cudnnDataType_t dataType;
+            cudnnTensorFormat_t format;
+            int nbDims;
+            int filterDimA[3];
+            checkCUDNN(cudnnGetFilterNdDescriptor(linLayerMatDesc, 3, &dataType,
+                                                     &format, &nbDims, filterDimA));
+            std::cout<<"Wgs Dims: "<<nbDims<<"  ("<<filterDimA[0]<<", "<<filterDimA[1]<<", "<<filterDimA[2]<<")\n";
+
+            // here we should fill the params data into linLayerMat
+
+            checkCUDNN(cudnnDestroyFilterDescriptor(linLayerMatDesc));
+
+            // get bias pointer
+            cudnnFilterDescriptor_t linLayerBiasDesc;
+            checkCUDNN(cudnnCreateFilterDescriptor(&linLayerBiasDesc));
+            float *linLayerBias;
+
+            checkCUDNN(cudnnGetRNNLinLayerBiasParams(net->cudnnHandle, rnnDesc,
+                    i, srcTensorDesc, paramDesc, paramsSpace,
+                    j, linLayerBiasDesc, (void **)&linLayerBias));
+
+            if(linLayerMat == nullptr) {
+                FatalError("LSTM No bias in hidden layer");
+            }
+
+            checkCUDNN(cudnnGetFilterNdDescriptor(linLayerBiasDesc, 3, &dataType,
+                                                     &format, &nbDims, filterDimA));
+            std::cout<<"bias Dims: "<<nbDims<<"  ("<<filterDimA[0]<<", "<<filterDimA[1]<<", "<<filterDimA[2]<<")\n";
+
+            // here we should fill the params data into linLayerBiasDesc
+
+            checkCUDNN(cudnnDestroyFilterDescriptor(linLayerBiasDesc));
+
+        }
+    }
+
+
+
+
+    checkCUDNN( cudnnCreateTensorDescriptor(&hiddenStateTensorDesc));
+    checkCUDNN( cudnnSetTensor4dDescriptor(hiddenStateTensorDesc,
+                                           net->tensorFormat, net->dataType, 2*n, c, h, w) );
+    checkCuda( cudaMalloc(&hiddenStateData, 2*input_dim.tot()*sizeof(dnnType)) );
+    checkCUDNN( cudnnCreateTensorDescriptor(&cellStateTensorDesc));
+    checkCUDNN( cudnnSetTensor4dDescriptor(cellStateTensorDesc,
+                                           net->tensorFormat, net->dataType, 2*n, c, h, w) );
+    checkCuda( cudaMalloc(&cellStateData, 2*input_dim.tot()*sizeof(dnnType)) );
+
+
+    output_dim = input_dim;
+    output_dim.c = hiddensize*2;
+    checkCUDNN( cudnnSetTensor4dDescriptor(dstTensorDesc,
+                                           net->tensorFormat, net->dataType, output_dim.n, output_dim.c, output_dim.h, output_dim.w) );
+
+
+
+
+    //allocate data for infer result
+    checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
+}
+
+LSTM::~LSTM() {
+
+    checkCuda( cudaFree(dstData) );
+}
+
+dnnType* LSTM::infer(dataDim_t &dim, dnnType* srcData) {
+
+    checkCUDNN(cudnnRNNForwardInference(
+        net->cudnnHandle, rnnDesc, 1,
+        &srcTensorDesc, srcData,
+        hiddenStateTensorDesc, hiddenStateData,
+        cellStateTensorDesc, cellStateData,
+        paramDesc, paramsSpace,
+        &dstTensorDesc, dstData,
+        hiddenStateTensorDesc, hiddenStateData,
+        cellStateTensorDesc, cellStateData,
+        workSpace, ws_sizeInBytes
+    ));
+
+    return dstData;
+}
+
+}}