Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix GPU/CPU Conv2DBackpropInputOp check error.
For empty `out_backprop` inputs (e.g. `[3, 1, 0, 1]`), the current
CPU/GPU kernels fail (one with dnnl, the other with cudnn). Added
a shortcut path to return a zero input.

Note that the XLA test currently fails due to an incorrect result
size bug (b/239598470), so it is currently disabled.

PiperOrigin-RevId: 462217998
  • Loading branch information
cantonios authored and tensorflower-gardener committed Jul 20, 2022
1 parent c55b476 commit 27a65a4
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 0 deletions.
19 changes: 19 additions & 0 deletions tensorflow/core/kernels/conv_grad_input_ops.h
Expand Up @@ -37,6 +37,7 @@ limitations under the License.
#include "tensorflow/core/kernels/conv_2d.h"
#include "tensorflow/core/kernels/conv_grad_ops.h"
#include "tensorflow/core/kernels/conv_grad_shape_utils.h"
#include "tensorflow/core/kernels/fill_functor.h"
#ifdef TENSORFLOW_USE_LIBXSMM_CONVOLUTIONS
#include "tensorflow/core/kernels/xsmm_conv2d.h"
#endif
Expand Down Expand Up @@ -436,6 +437,15 @@ class Conv2DBackpropInputOp : public OpKernel {
return;
}

// If shapes are valid but `out_backprop` is empty, in_backprop should be
// set to all zeros. Otherwise, cudnn/dnnl fail with an empty input.
if (out_backprop.NumElements() == 0) {
functor::SetZeroFunctor<Device, T> set_zero;
set_zero(context->eigen_device<Device>(),
in_backprop->template flat<T>());
return;
}

// For now we take the stride from the second and third dimensions only (we
// do not support striding on the batch or depth dimension).
const int stride_rows = GetTensorDim(strides_, data_format_, 'H');
Expand Down Expand Up @@ -554,6 +564,15 @@ class Conv2DCustomBackpropInputOp : public OpKernel {
return;
}

// If shapes are valid but `out_backprop` is empty, in_backprop should be
// set to all zeros. Otherwise, cudnn/dnnl fail with an empty input.
if (out_backprop.NumElements() == 0) {
functor::SetZeroFunctor<Device, T> set_zero;
set_zero(context->eigen_device<Device>(),
in_backprop->template flat<T>());
return;
}

// TODO(ezhulenev): Remove custom kernel and move XSMM support to
// LaunchConv2DBackpropInputOp functor.
#if defined TENSORFLOW_USE_LIBXSMM_CONVOLUTIONS && \
Expand Down
17 changes: 17 additions & 0 deletions tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
Expand Up @@ -1103,6 +1103,23 @@ def testConv2DInputSizesContainsOnlySpatialDimensionsBackpropInput(self):
use_gpu=use_gpu,
err=1e-5)

@test_util.run_in_graph_and_eager_modes
@test_util.disable_xla("b/239598470")
def testConv2DBackpropInputDegenerateBackpropInput(self):
input_sizes = [3, 1, 1, 2]
expected_output = np.zeros(input_sizes).flatten()
for (data_format, use_gpu) in GetTestConfigs():
self._RunAndVerifyBackpropInput(
input_sizes=input_sizes,
filter_sizes=[1, 3, 2, 3],
output_sizes=[3, 1, 0, 3],
strides=[1, 2],
padding="VALID",
expected=expected_output,
data_format=data_format,
use_gpu=use_gpu,
err=1e-5)

# Testing for backprops
def _RunAndVerifyBackpropFilter(self,
input_sizes,
Expand Down

0 comments on commit 27a65a4

Please sign in to comment.