Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix conv2d crash when input size is empty.
If the input is empty (so convolution is only applied to padding), and if
the filter and padding sizes are still valid, then the output will be
all-zeros.  This previously caused a division-by-zero crash in multiple
kernels.

PiperOrigin-RevId: 463179650
  • Loading branch information
cantonios authored and tensorflower-gardener committed Jul 25, 2022
1 parent de2e1fe commit 611d80d
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 0 deletions.
10 changes: 10 additions & 0 deletions tensorflow/core/kernels/conv_ops.cc
Expand Up @@ -44,6 +44,7 @@ limitations under the License.
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/kernels/conv_2d.h"
#include "tensorflow/core/kernels/deep_conv2d.h"
#include "tensorflow/core/kernels/fill_functor.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
Expand Down Expand Up @@ -701,6 +702,15 @@ class Conv2DOp : public BinaryOp<T> {
return;
}

// If the input is empty, result can only be due to padding.
if (input.NumElements() == 0) {
// Zero-out output and return.
functor::SetZeroFunctor<Device, T>()(context->eigen_device<Device>(),
output->template flat<T>());

return;
}

#ifdef TENSORFLOW_USE_LIBXSMM_CONVOLUTIONS
if (params_.padding != EXPLICIT &&
LaunchXsmmConvOp<Device, T>::Run(
Expand Down
9 changes: 9 additions & 0 deletions tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
Expand Up @@ -759,6 +759,15 @@ def testConv2DExplicitPaddingWithDilations(self):
padding=[[2, 1], [1, 2]],
dilations=[2, 3])

@test_util.run_in_graph_and_eager_modes()
def testConv2dOnlyPaddingReturnsZeros(self):
self._VerifyValues(
tensor_in_sizes=[1, 0, 2, 1],
filter_in_sizes=[1, 1, 1, 1],
strides=[1, 1],
padding=[[1, 1], [1, 1]],
expected=[0, 0, 0, 0, 0, 0, 0, 0])

def testConv2DExplicitPaddingWithLayoutOptimizer(self):
# Test with Grappler's layout optimizer, to ensure the layout optimizer
# handles explicit padding correctly.
Expand Down

0 comments on commit 611d80d

Please sign in to comment.