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')
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,20 +31,25 @@ 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 }));
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;
}
}
// 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],
input_desc,
filter_desc,
......@@ -167,20 +172,25 @@ 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 }));
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;
}
}
// 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],
output_desc,
input_desc,
......
Dragon/src/operators/vision/cudnn_conv2d_transpose_op.cc
......@@ -31,36 +31,42 @@ 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;
}
}
CUDNN_CHECK(cudnnGetConvolutionBackwardDataAlgorithm(handle[0],
filter_desc,
input_desc,
conv_desc,
filter_desc,
input_desc,
conv_desc,
output_desc,
CUDNN_CONVOLUTION_BWD_DATA_SPECIFY_WORKSPACE_LIMIT,
WORKSPACE_LIMIT_BYTES,
WORKSPACE_LIMIT_BYTES,
&fwd_algo));
CUDNN_CHECK(cudnnGetConvolutionBackwardDataWorkspaceSize(handle[0],
filter_desc,
input_desc,
conv_desc,
filter_desc,
input_desc,
conv_desc,
output_desc,
fwd_algo,
fwd_algo,
&workspace_fwd_data_size));
Tensor* buffer = ws()->GetBuffer();
......@@ -75,7 +81,7 @@ void CuDNNConv2dTransposeOp<Context>::RunWithType() {
for (int g = 0; g < cudnn_group; g++) {
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,
input_desc, Xdata + this->x_offset * g,
conv_desc,
......@@ -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;
}
......@@ -318,4 +330,4 @@ DEPLOY_CUDNN(Conv2dTransposeGradient);
} // namespace dragon
#endif // WITH_CUDNN
\ No newline at end of file
#endif // WITH_CUDNN
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;
......
Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!