Skip to content

SeetaResearch / Dragon

Go to a project
Commit 387a3675 by Ting PAN
Options

Fix the potential crash on cuDNNGroupConvolution

1 parent 3d2abe69
Inline Side-by-side
Showing with 38 additions and 16 deletions
Dragon/python/setup.py
...@@ -36,7 +36,7 @@ find_packages('dragon') ...@@ -36,7 +36,7 @@ find_packages('dragon')
find_modules() find_modules()
setup(name = 'dragon', setup(name = 'dragon',
version='0.2.1.9', version='0.2.1.10',
description = 'Dragon: A Computation Graph Virtual Machine Based Deep Learning Framework', description = 'Dragon: A Computation Graph Virtual Machine Based Deep Learning Framework',
url='https://github.com/neopenx/Dragon', url='https://github.com/neopenx/Dragon',
author='Ting Pan', author='Ting Pan',
......
Dragon/src/operators/vision/cudnn_conv2d_op.cc
...@@ -31,19 +31,24 @@ void CuDNNConv2dOp<Context>::RunWithType() { ...@@ -31,19 +31,24 @@ void CuDNNConv2dOp<Context>::RunWithType() {
cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(0).dims(), cudnn_group); cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(0).dims(), cudnn_group);
cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, output(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()) { if (HasBias()) {
bias_offset = this->num_output / cudnn_group; bias_offset = this->num_output / cudnn_group;
if (this->data_format == "NCHW") { if (this->data_format == "NCHW") {
cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 })); 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->x_offset = input(0).count(1) / cudnn_group;
this->y_offset = output(0)->count(1) / cudnn_group; this->y_offset = output(0)->count(1) / cudnn_group;
} else if (this->data_format == "NHWC") { } 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->x_offset = input(0).dim(-1) / cudnn_group;
this->y_offset = output(0)->dim(-1) / cudnn_group; this->y_offset = output(0)->dim(-1) / cudnn_group;
} }
}
CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm(handle[0], CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm(handle[0],
input_desc, input_desc,
...@@ -167,19 +172,24 @@ void CuDNNConv2dGradientOp<Context>::RunWithType() { ...@@ -167,19 +172,24 @@ void CuDNNConv2dGradientOp<Context>::RunWithType() {
cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(-1).dims(), cudnn_group); cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(-1).dims(), cudnn_group);
cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, input(0).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()) { if (HasBias()) {
bias_offset = this->num_output / cudnn_group; bias_offset = this->num_output / cudnn_group;
if (this->data_format == "NCHW") { if (this->data_format == "NCHW") {
cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 })); 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->x_offset = input(0).count(1) / cudnn_group;
this->y_offset = input(-1).count(1) / cudnn_group; this->y_offset = input(-1).count(1) / cudnn_group;
} else if (this->data_format == "NHWC") { } 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->x_offset = input(0).dim(-1) / cudnn_group;
this->y_offset = input(-1).dim(-1) / cudnn_group; this->y_offset = input(-1).dim(-1) / cudnn_group;
} }
}
CUDNN_CHECK(cudnnGetConvolutionBackwardFilterAlgorithm(handle[0], CUDNN_CHECK(cudnnGetConvolutionBackwardFilterAlgorithm(handle[0],
output_desc, output_desc,
......
Dragon/src/operators/vision/cudnn_conv2d_transpose_op.cc
...@@ -31,16 +31,22 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() { ...@@ -31,16 +31,22 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() {
cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(0).dims(), cudnn_group); cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(0).dims(), cudnn_group);
cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, output(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()) { if (HasBias()) {
bias_offset = this->num_output / cudnn_group; bias_offset = this->num_output / cudnn_group;
if (this->data_format == "NCHW") { if (this->data_format == "NCHW") {
cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 })); 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->x_offset = input(0).count(1) / cudnn_group;
this->y_offset = output(0)->count(1) / cudnn_group; this->y_offset = output(0)->count(1) / cudnn_group;
} } else if (this->data_format == "NHWC") {
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->x_offset = input(0).dim(-1) / cudnn_group;
this->y_offset = output(0)->dim(-1) / cudnn_group; this->y_offset = output(0)->dim(-1) / cudnn_group;
} }
...@@ -169,16 +175,22 @@ void CuDNNConv2dTransposeGradientOp<Context>::RunWithType() { ...@@ -169,16 +175,22 @@ void CuDNNConv2dTransposeGradientOp<Context>::RunWithType() {
cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(-1).dims(), cudnn_group); cudnnSetTensor4dDescWithGroup<T>(&input_desc, this->data_format, input(-1).dims(), cudnn_group);
cudnnSetTensor4dDescWithGroup<T>(&output_desc, this->data_format, input(0).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()) { if (HasBias()) {
bias_offset = this->num_output / cudnn_group; bias_offset = this->num_output / cudnn_group;
if (this->data_format == "NCHW") { if (this->data_format == "NCHW") {
cudnnSetTensor4dDesc<T>(&bias_desc, this->data_format, vector<TIndex>({ 1, bias_offset, 1, 1 })); 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->x_offset = input(0).count(1) / cudnn_group;
this->y_offset = input(-1).count(1) / cudnn_group; this->y_offset = input(-1).count(1) / cudnn_group;
} } else if (this->data_format == "NHWC") {
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->x_offset = input(0).dim(-1) / cudnn_group;
this->y_offset = input(-1).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, ...@@ -3935,7 +3935,7 @@ __global__ void _ROIPoolingGrad(const int count,
if (!in_roi) continue; 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 T* offset_dy = dy + y_offset;
const int* offset_mask = mask + y_offset; const int* offset_mask = mask + y_offset;
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!