Commit 387a3675 by Ting PAN

Fix the potential crash on cuDNNGroupConvolution

1 parent 3d2abe69
......@@ -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',
......
......@@ -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,
......
......@@ -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;
}
......
......@@ -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!