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 52 additions and 31 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,20 +31,25 @@ void CuDNNConv2dOp<Context>::RunWithType() { ...@@ -31,20 +31,25 @@ 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 }));
this->x_offset = input(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 })); 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;
} }
} }
// 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") {
this->x_offset = input(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,
filter_desc, filter_desc,
...@@ -167,20 +172,25 @@ void CuDNNConv2dGradientOp<Context>::RunWithType() { ...@@ -167,20 +172,25 @@ 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 }));
this->x_offset = input(0).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 })); 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;
} }
} }
// 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") {
this->x_offset = input(0).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,
input_desc, input_desc,
......
Dragon/src/operators/vision/cudnn_conv2d_transpose_op.cc
...@@ -31,36 +31,42 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() { ...@@ -31,36 +31,42 @@ 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;
} }
} }
CUDNN_CHECK(cudnnGetConvolutionBackwardDataAlgorithm(handle[0], CUDNN_CHECK(cudnnGetConvolutionBackwardDataAlgorithm(handle[0],
filter_desc, filter_desc,
input_desc, input_desc,
conv_desc, conv_desc,
output_desc, output_desc,
CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT, CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT,
WORKSPACE_LIMIT_BYTES, WORKSPACE_LIMIT_BYTES,
&fwd_algo)); &fwd_algo));
CUDNN_CHECK(cudnnGetConvolutionBackwardDataWorkspaceSize(handle[0], CUDNN_CHECK(cudnnGetConvolutionBackwardDataWorkspaceSize(handle[0],
filter_desc, filter_desc,
input_desc, input_desc,
conv_desc, conv_desc,
output_desc, output_desc,
fwd_algo, fwd_algo,
&workspace_fwd_data_size)); &workspace_fwd_data_size));
Tensor* buffer = ws()->GetBuffer(); Tensor* buffer = ws()->GetBuffer();
...@@ -75,7 +81,7 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() { ...@@ -75,7 +81,7 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() {
for (int g = 0; g < cudnn_group; g++) { for (int g = 0; g < cudnn_group; g++) {
auto* workspace = buffer->template mutable_data<char, Context>(); auto* workspace = buffer->template mutable_data<char, Context>();
CUDNN_CHECK(cudnnConvolutionBackwardData(handle[g], CUDNN_CHECK(cudnnConvolutionBackwardData(handle[g],
CUDNNType<T>::one, filter_desc, Wdata + this->weight_offset * g, CUDNNType<T>::one, filter_desc, Wdata + this->weight_offset * g,
input_desc, Xdata + this->x_offset * g, input_desc, Xdata + this->x_offset * g,
conv_desc, conv_desc,
...@@ -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;
} }
...@@ -318,4 +330,4 @@ DEPLOY_CUDNN(Conv2dTransposeGradient); ...@@ -318,4 +330,4 @@ DEPLOY_CUDNN(Conv2dTransposeGradient);
} // namespace dragon } // namespace dragon
#endif // WITH_CUDNN #endif // WITH_CUDNN
\ No newline at end of file
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!