deconv layer cudnn

2019-10-30 00:05:22 +01:00
parent 02936fa928
commit f9afee2f3b
7 changed files with 197 additions and 102 deletions
@@ -4,9 +4,121 @@

 namespace tk { namespace dnn {

+void Conv2d::initCUDNN(bool back) {
+
+    cudnnTensorDescriptor_t srcTensor = srcTensorDesc;
+    cudnnTensorDescriptor_t dstTensor = dstTensorDesc;
+
+    dataDim_t idim, odim;
+    if(!back) {
+        idim = input_dim;
+        odim = output_dim;
+    } else {
+        idim = output_dim;
+        odim = input_dim;
+    }
+    idim.print();
+    odim.print();
+
+    checkCUDNN( cudnnCreateFilterDescriptor(&filterDesc) );
+    checkCUDNN( cudnnCreateConvolutionDescriptor(&convDesc) );
+    checkCUDNN( cudnnCreateTensorDescriptor(&biasTensorDesc) );
+
+    // input tensor dim
+    checkCUDNN( cudnnSetTensor4dDescriptor(srcTensor,
+                                           net->tensorFormat, net->dataType, idim.n, idim.c, idim.h, idim.w) );
+
+    checkCUDNN( cudnnSetFilter4dDescriptor(filterDesc,
+                                           net->dataType, net->tensorFormat, odim.c, idim.c,
+                                           kernelH, kernelW) );
+
+    checkCUDNN( cudnnSetConvolution2dDescriptor(convDesc,
+                                                paddingH, paddingW, // padding
+                                                strideH, strideW, // stride
+                                                1,1, // upscale
+                                                CUDNN_CROSS_CORRELATION, CUDNN_DATA_FLOAT) );
+
+    // check dimension of convolution output
+    dataDim_t tmpdim;
+    checkCUDNN( cudnnGetConvolution2dForwardOutputDim(
+            convDesc, srcTensor, filterDesc,
+            &tmpdim.n, &tmpdim.c, &tmpdim.h, &tmpdim.w) );
+    if(odim.n != tmpdim.n || odim.c != tmpdim.c || odim.h != tmpdim.h || odim.w != tmpdim.w) {
+        std::cout<<"tkdim: "; odim.print();
+        std::cout<<"cudnndim: "; tmpdim.print();
+        FatalError("Eror conv dimension mismatch");
+    }
+
+    checkCUDNN( cudnnSetTensor4dDescriptor(dstTensor,
+                                           net->tensorFormat, net->dataType, odim.n, odim.c, odim.h, odim.w) );
+
+    checkCUDNN( cudnnSetTensor4dDescriptor(biasTensorDesc,
+                                           net->tensorFormat, net->dataType,
+                                           1, output_dim.c, 1, 1) );
+
+    // init workspace
+    workSpace = NULL;
+    ws_sizeInBytes = 0;
+    if(back) {
+        checkCUDNN( cudnnGetConvolutionBackwardDataAlgorithm(net->cudnnHandle,
+                                                        filterDesc, dstTensor, convDesc, srcTensor,
+                                                        CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST, 0, &bwAlgo) );
+        checkCUDNN(cudnnGetConvolutionBackwardDataWorkspaceSize(net->cudnnHandle,
+                                                           filterDesc, dstTensor, convDesc, srcTensor,
+                                                           bwAlgo, &ws_sizeInBytes));
+
+        // invert tensors
+        srcTensorDesc = dstTensor;
+        dstTensorDesc = srcTensor;
+    } else {
+        checkCUDNN( cudnnGetConvolutionForwardAlgorithm(net->cudnnHandle,
+                                                        srcTensor, filterDesc, convDesc, dstTensor,
+                                                        CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &fwAlgo) );
+        checkCUDNN(cudnnGetConvolutionForwardWorkspaceSize(net->cudnnHandle,
+                                                           srcTensor, filterDesc, convDesc, dstTensor,
+                                                           fwAlgo, &ws_sizeInBytes));
+    }
+}
+
+void Conv2d::inferCUDNN(dnnType* srcData, bool back) {
+
+    dnnType alpha = dnnType(1);
+    dnnType beta  = dnnType(0);
+    if(back) {
+        checkCUDNN(cudnnConvolutionBackwardData(net->cudnnHandle,
+                                                &alpha, filterDesc, data_d,
+                                                srcTensorDesc, srcData,
+                                                convDesc, bwAlgo, workSpace, ws_sizeInBytes,
+                                                &beta, dstTensorDesc, dstData));
+    } else {
+        checkCUDNN(cudnnConvolutionForward(net->cudnnHandle,
+                                           &alpha, srcTensorDesc, srcData, filterDesc,
+                                           data_d, convDesc, fwAlgo, workSpace, ws_sizeInBytes,
+                                           &beta, dstTensorDesc, dstData));
+    }
+
+    if(!batchnorm) {
+        // bias
+        alpha = dnnType(1);
+        beta  = dnnType(0);
+        checkCUDNN( cudnnAddTensor(net->cudnnHandle,
+                                   &alpha, biasTensorDesc, bias_d,
+                                   &beta, dstTensorDesc, dstData) );
+    } else {
+        alpha = dnnType(1);
+        beta  = dnnType(0);
+        cudnnBatchNormalizationForwardInference(net->cudnnHandle,
+                                                CUDNN_BATCHNORM_SPATIAL, &alpha, &beta,
+                                                dstTensorDesc, dstData, dstTensorDesc,
+                                                dstData, biasTensorDesc, //same tensor descriptor as bias
+                                                scales_d, bias_d, mean_d, variance_d,
+                                                CUDNN_BN_MIN_EPSILON);
+    }
+}
+
 Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW, 
                int strideH, int strideW, int paddingH, int paddingW,
-                std::string fname_weights, bool batchnorm) : 
+                std::string fname_weights, bool batchnorm, bool deConv) :
    
    LayerWgs(net, net->getOutputDim().c, out_ch, kernelH, kernelW, 1, 
             fname_weights, batchnorm) {
@@ -17,64 +129,28 @@ Conv2d::Conv2d( Network *net, int out_ch, int kernelH, int kernelW,
    this->strideW = strideW;
    this->paddingH = paddingH;
    this->paddingW = paddingW;
+    this->deConv = deConv;

-    checkCUDNN( cudnnCreateFilterDescriptor(&filterDesc) );
-    checkCUDNN( cudnnCreateConvolutionDescriptor(&convDesc) );
-    checkCUDNN( cudnnCreateTensorDescriptor(&biasTensorDesc) );
-
-    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, c, h, w) );
-
-    checkCUDNN( cudnnSetFilter4dDescriptor(filterDesc,
-                net->dataType, net->tensorFormat, out_ch, input_dim.c, 
-                kernelH, kernelW) );
-
-    checkCUDNN( cudnnSetConvolution2dDescriptor(convDesc,
-                paddingH, paddingW, // padding
-                strideH, strideW, // stride
-                1,1, // upscale
-                CUDNN_CROSS_CORRELATION, CUDNN_DATA_FLOAT) );
-
-    // find dimension of convolution output
-    checkCUDNN( cudnnGetConvolution2dForwardOutputDim(
-                convDesc, srcTensorDesc, filterDesc,
-                &n, &c, &h, &w) );
-
-    checkCUDNN( cudnnSetTensor4dDescriptor(dstTensorDesc,
-                net->tensorFormat, net->dataType, n, c, h, w) );
-                
-    checkCUDNN( cudnnGetConvolutionForwardAlgorithm(net->cudnnHandle,
-                srcTensorDesc, filterDesc, convDesc, dstTensorDesc,
-                CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo) );
- 
-    workSpace = NULL;
-    ws_sizeInBytes = 0;
-
-    checkCUDNN( cudnnGetConvolutionForwardWorkspaceSize(net->cudnnHandle,
-                srcTensorDesc, filterDesc, convDesc, dstTensorDesc,
-                algo, &ws_sizeInBytes) );
+    if(!deConv) {
+        output_dim.n = input_dim.n;
+        output_dim.c = out_ch;
+        output_dim.h = (input_dim.h + 2 * paddingH - kernelH) / strideH + 1;
+        output_dim.w = (input_dim.w + 2 * paddingW - kernelW) / strideW + 1;
+        output_dim.l = 1;
+    } else {
+        output_dim.n = input_dim.n;
+        output_dim.c = out_ch;
+        output_dim.h = (input_dim.h * strideH) - 2*paddingH + kernelH -1;
+        output_dim.w = (input_dim.w * strideW) - 2*paddingW + kernelW -1;
+        output_dim.l = 1;
+    }
+    initCUDNN(deConv);

+    // allocate warkspace
    if (ws_sizeInBytes!=0) {
        checkCuda( cudaMalloc(&workSpace, ws_sizeInBytes) );
    }

-
-    checkCUDNN( cudnnSetTensor4dDescriptor(biasTensorDesc,
-                net->tensorFormat, net->dataType,
-                1, out_ch, 1, 1) );
-
-
-    output_dim.n = n;
-    output_dim.c = c;
-    output_dim.h = h;
-    output_dim.w = w;
-    output_dim.l = 1;
-
    //allocate data for infer result
    checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
 }
@@ -93,35 +169,25 @@ Conv2d::~Conv2d() {

 dnnType* Conv2d::infer(dataDim_t &dim, dnnType* srcData) {

+    if(deConv) {
+        FatalError("you must use DeConv class for Deconvolutional layers");
+    }

    // convolution
-    dnnType alpha = dnnType(1);
-    dnnType beta  = dnnType(0);
-    checkCUDNN( cudnnConvolutionForward(net->cudnnHandle,
-                &alpha, srcTensorDesc, srcData, filterDesc,
-                data_d, convDesc, algo, workSpace, ws_sizeInBytes,
-                &beta, dstTensorDesc, dstData) );
+    inferCUDNN(srcData, false);

-    if(!batchnorm) {
-        // bias
-        alpha = dnnType(1);
-        beta  = dnnType(1);
-        checkCUDNN( cudnnAddTensor(net->cudnnHandle,
-                    &alpha, biasTensorDesc, bias_d,
-                    &beta, dstTensorDesc, dstData) );
-    } else {
-        float one = 1;
-        float zero = 0;
-        cudnnBatchNormalizationForwardInference(net->cudnnHandle,
-            CUDNN_BATCHNORM_SPATIAL, &one, &zero, 
-            dstTensorDesc, dstData, dstTensorDesc, 
-            dstData, biasTensorDesc, //same tensor descriptor as bias 
-            scales_d, bias_d, mean_d, variance_d, 
-            CUDNN_BN_MIN_EPSILON);
-    }
    //update data dimensions    
    dim = output_dim;
+    return dstData;
+}

+dnnType* DeConv2d::infer(dataDim_t &dim, dnnType* srcData) {
+
+    // convolution
+    inferCUDNN(srcData, true);
+
+    //update data dimensions
+    dim = output_dim;
    return dstData;
 }

@@ -16,18 +16,18 @@ LayerWgs::LayerWgs(Network *net, int inputs, int outputs,
   
    std::cout<<"Reading weights: I="<<inputs<<" O="<<outputs<<" KERNEL="<<kh<<"x"<<kw<<"x"<<kl<<"\n";
    int seek = 0;
-    readBinaryFile(weights_path.c_str(), inputs*outputs*kh*kw*kl, &data_h, &data_d, seek);
+    readBinaryFile(weights_path.c_str(), inputs*outputs*kh*kw*kl, &data_h, &data_d, seek, net->dontLoadWeights);
    seek += inputs*outputs*kh*kw*kl;
-    readBinaryFile(weights_path.c_str(), outputs, &bias_h, &bias_d, seek);
+    readBinaryFile(weights_path.c_str(), outputs, &bias_h, &bias_d, seek, net->dontLoadWeights);

    this->batchnorm = batchnorm;
    if(batchnorm) {
        seek += outputs;
-        readBinaryFile(weights_path.c_str(), outputs, &scales_h, &scales_d, seek);
+        readBinaryFile(weights_path.c_str(), outputs, &scales_h, &scales_d, seek, net->dontLoadWeights);
        seek += outputs;
-        readBinaryFile(weights_path.c_str(), outputs, &mean_h, &mean_d, seek);
+        readBinaryFile(weights_path.c_str(), outputs, &mean_h, &mean_d, seek, net->dontLoadWeights);
        seek += outputs;
-        readBinaryFile(weights_path.c_str(), outputs, &variance_h, &variance_d, seek);
+        readBinaryFile(weights_path.c_str(), outputs, &variance_h, &variance_d, seek, net->dontLoadWeights);

        float eps = CUDNN_BN_MIN_EPSILON;

@@ -17,6 +17,7 @@ Network::Network(dataDim_t input_dim) {
             <<", CUDNN v"<<cu_ver<<")\n";
    dataType = CUDNN_DATA_FLOAT;
    tensorFormat = CUDNN_TENSOR_NCHW;
+    dontLoadWeights = false;

    checkCUDNN( cudnnCreate(&cudnnHandle) );
    checkERROR( cublasCreate(&cublasHandle) );
@@ -21,26 +21,29 @@ bool fileExist(const char *fname) {
 }


-void readBinaryFile(std::string fname, int size, dnnType** data_h, dnnType** data_d, int seek)
+void readBinaryFile(std::string fname, int size, dnnType** data_h, dnnType** data_d, int seek, bool skipLoad)
 {
-    std::ifstream dataFile (fname, std::ios::in | std::ios::binary);
-    std::stringstream error_s;
-    if (!dataFile)
-    {
-        error_s << "Error opening file " << fname; 
-        FatalError(error_s.str());
-    }
-
-    if(seek != 0) {
-        dataFile.seekg(seek*sizeof(dnnType), dataFile.cur);
-    }
-
    int size_b = size*sizeof(dnnType);
    *data_h = new dnnType[size];
-    if (!dataFile.read ((char*) *data_h, size_b)) 
-    {
-        error_s << "Error reading file " << fname; 
-        FatalError(error_s.str());
+
+    if(!skipLoad) {
+        std::ifstream dataFile(fname, std::ios::in | std::ios::binary);
+        std::stringstream error_s;
+        if (!dataFile) {
+            error_s << "Error opening file " << fname;
+            FatalError(error_s.str());
+        }
+
+        if (seek != 0) {
+            dataFile.seekg(seek * sizeof(dnnType), dataFile.cur);
+        }
+
+        if (!dataFile.read((char *) *data_h, size_b)) {
+            error_s << "Error reading file " << fname;
+            FatalError(error_s.str());
+        }
+    } else {
+        std::cout<<COL_RED<<"WARNING: skipping data load, this should only used in debug\n"<<COL_END;
    }
    
    checkCuda( cudaMalloc(data_d, size_b) );