Fix the potential crash on cuDNNGroupConvolution

Dragon/python/setup.py

@@ -36,7 +36,7 @@ find_packages('dragon')
 find_modules()
 
 setup(name = 'dragon',
-      version='0.2.1.9',
+      version='0.2.1.10',
       description = 'Dragon: A Computation Graph Virtual Machine Based Deep Learning Framework',
       url='https://github.com/neopenx/Dragon',
       author='Ting Pan',

Dragon/src/operators/vision/cudnn_conv2d_op.cc

@@ -31,19 +31,24 @@ void CuDNNConv2dOp<Context>::RunWithType() {
     cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(0).dims(), cudnn_group);
     cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, output(0)->dims(), cudnn_group);
 
-    //  determine the bias shape and misc
+    //  determine the bias shape
     if (HasBias()) {
         bias_offset = this->num_output / cudnn_group;
         if (this->data_format == "NCHW") {
             cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 }));
+        } else if (this->data_format == "NHWC") {
+            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
+        }
+    }
+
+    // determine the misc
+    if (this->data_format == "NCHW") {
         this->x_offset = input(0).count(1) / cudnn_group;
         this->y_offset = output(0)->count(1) / cudnn_group;
     }  else if (this->data_format == "NHWC") {
-            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
         this->x_offset = input(0).dim(-1) / cudnn_group;
         this->y_offset = output(0)->dim(-1) / cudnn_group;
     }
-    }
 
     CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm(handle[0],
                                                    input_desc,
@@ -167,19 +172,24 @@ void CuDNNConv2dGradientOp<Context>::RunWithType() {
     cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(-1).dims(), cudnn_group);
     cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, input(0).dims(), cudnn_group);
 
-    //  determine the bias shape and misc
+    //  determine the bias shape
     if (HasBias()) {
         bias_offset = this->num_output / cudnn_group;
         if (this->data_format == "NCHW") {
             cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 }));
+        } else if (this->data_format == "NHWC") {
+            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
+        }
+    }
+
+    // determine the misc
+    if (this->data_format == "NCHW") {
         this->x_offset = input(0).count(1) / cudnn_group;
         this->y_offset = input(-1).count(1) / cudnn_group;
     } else if (this->data_format == "NHWC") {
-            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
         this->x_offset = input(0).dim(-1) / cudnn_group;
         this->y_offset = input(-1).dim(-1) / cudnn_group;
     }
-    }
 
     CUDNN_CHECK(cudnnGetConvolutionBackwardFilterAlgorithm(handle[0],
                                                          output_desc,

Dragon/src/operators/vision/cudnn_conv2d_transpose_op.cc

@@ -31,16 +31,22 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() {
     cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(0).dims(), cudnn_group);
     cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, output(0)->dims(), cudnn_group);
 
-    //  determine the bias shape and misc
+    //  determine the bias shape
     if (HasBias()) {
         bias_offset = this->num_output / cudnn_group;
         if (this->data_format == "NCHW") {
             cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 }));
+        } else if (this->data_format == "NHWC") {
+            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
+        }
+    }
+
+    //  determine the misc
+    if (HasBias()) {
+        if (this->data_format == "NCHW") {
             this->x_offset = input(0).count(1) / cudnn_group;
             this->y_offset = output(0)->count(1) / cudnn_group;
-        }
-        else if (this->data_format == "NHWC") {
-            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
+        } else if (this->data_format == "NHWC") {
             this->x_offset = input(0).dim(-1) / cudnn_group;
             this->y_offset = output(0)->dim(-1) / cudnn_group;
         }
@@ -169,16 +175,22 @@ void CuDNNConv2dTransposeGradientOp<Context>::RunWithType() {
     cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(-1).dims(), cudnn_group);
     cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, input(0).dims(), cudnn_group);
 
-    //  determine the bias shape and misc
+    //  determine the bias shape
     if (HasBias()) {
         bias_offset = this->num_output / cudnn_group;
         if (this->data_format == "NCHW") {
             cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 }));
+        } else if (this->data_format == "NHWC") {
+            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
+        }
+    }
+
+    //  determine the misc
+    if (HasBias()) {
+        if (this->data_format == "NCHW") {
             this->x_offset = input(0).count(1) / cudnn_group;
             this->y_offset = input(-1).count(1) / cudnn_group;
-        }
-        else if (this->data_format == "NHWC") {
-            cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, 1, 1, bias_offset }));
+        } else if (this->data_format == "NHWC") {
             this->x_offset = input(0).dim(-1) / cudnn_group;
             this->y_offset = input(-1).dim(-1) / cudnn_group;
         }

Dragon/src/utils/op_kernel.cu

@@ -3935,7 +3935,7 @@ __global__ void _ROIPoolingGrad(const int count,
 
             if (!in_roi) continue;
 
-            int y_offset = (n * channels + c) * pool_h * pool_w;
+            int y_offset = (roi_n * channels + c) * pool_h * pool_w;
             const T* offset_dy = dy + y_offset;
             const int* offset_mask = mask + y_offset;