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;
 }