102 lines
4.5 KiB
Diff
102 lines
4.5 KiB
Diff
From 311403edbc9816df80274bd1ea8b3c0c0f22c3fa Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Mon, 19 Apr 2021 16:00:40 -0700
|
|
Subject: [PATCH] Eliminate a division by 0 in 3D convolutions.
|
|
|
|
Also prevent a CHECK failed introduced in the most recent change.
|
|
|
|
PiperOrigin-RevId: 369322073
|
|
Change-Id: I4f609c028f89565fb2b49c3fdd20b63496582bae
|
|
---
|
|
tensorflow/core/kernels/conv_grad_ops_3d.cc | 42 +++++++++++++++++++++
|
|
1 file changed, 42 insertions(+)
|
|
|
|
diff --git a/tensorflow/core/kernels/conv_grad_ops_3d.cc b/tensorflow/core/kernels/conv_grad_ops_3d.cc
|
|
index 8c72d01578d6d..c0b57a7ae5606 100644
|
|
--- a/tensorflow/core/kernels/conv_grad_ops_3d.cc
|
|
+++ b/tensorflow/core/kernels/conv_grad_ops_3d.cc
|
|
@@ -239,6 +239,14 @@ class Conv3DBackpropInputOp : public OpKernel {
|
|
input_shape = context->input(0).shape();
|
|
}
|
|
|
|
+ OP_REQUIRES(context, input_shape.dims() == 5,
|
|
+ errors::InvalidArgument("input tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, filter_shape.dims() == 5,
|
|
+ errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, out_backprop_shape.dims() == 5,
|
|
+ errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
|
|
OP_REQUIRES(
|
|
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
|
|
errors::InvalidArgument("input and filter_sizes must have the same "
|
|
@@ -360,6 +368,14 @@ class Conv3DCustomBackpropInputOp : public OpKernel {
|
|
input_shape = context->input(0).shape();
|
|
}
|
|
|
|
+ OP_REQUIRES(context, input_shape.dims() == 5,
|
|
+ errors::InvalidArgument("input tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, filter_shape.dims() == 5,
|
|
+ errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, out_backprop_shape.dims() == 5,
|
|
+ errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
|
|
OP_REQUIRES(
|
|
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
|
|
errors::InvalidArgument("input and filter_sizes must have the same "
|
|
@@ -444,6 +460,11 @@ class Conv3DCustomBackpropInputOp : public OpKernel {
|
|
// contraction compared to sharding and matmuls.
|
|
const bool use_parallel_contraction = dims.batch_size == 1;
|
|
|
|
+ OP_REQUIRES(
|
|
+ context, work_unit_size > 0,
|
|
+ errors::InvalidArgument("input, filter_sizes and out_backprop tensors "
|
|
+ "must all have at least 1 element"));
|
|
+
|
|
const size_t shard_size =
|
|
use_parallel_contraction
|
|
? 1
|
|
@@ -724,6 +745,14 @@ class Conv3DBackpropFilterOp : public OpKernel {
|
|
filter_shape = context->input(1).shape();
|
|
}
|
|
|
|
+ OP_REQUIRES(context, input_shape.dims() == 5,
|
|
+ errors::InvalidArgument("input tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, filter_shape.dims() == 5,
|
|
+ errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, out_backprop_shape.dims() == 5,
|
|
+ errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
|
|
OP_REQUIRES(
|
|
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
|
|
errors::InvalidArgument("input and filter_sizes must have the same "
|
|
@@ -850,6 +879,14 @@ class Conv3DCustomBackpropFilterOp : public OpKernel {
|
|
filter_shape = context->input(1).shape();
|
|
}
|
|
|
|
+ OP_REQUIRES(context, input_shape.dims() == 5,
|
|
+ errors::InvalidArgument("input tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, filter_shape.dims() == 5,
|
|
+ errors::InvalidArgument("filter_sizes tensor must have 5 dimensions"));
|
|
+ OP_REQUIRES(
|
|
+ context, out_backprop_shape.dims() == 5,
|
|
+ errors::InvalidArgument("out_backprop tensor must have 5 dimensions"));
|
|
OP_REQUIRES(
|
|
context, input_shape.dim_size(4) == filter_shape.dim_size(3),
|
|
errors::InvalidArgument("input and filter_sizes must have the same "
|
|
@@ -936,6 +973,11 @@ class Conv3DCustomBackpropFilterOp : public OpKernel {
|
|
|
|
const int64 work_unit_size = size_A + size_B + size_C;
|
|
|
|
+ OP_REQUIRES(
|
|
+ context, work_unit_size > 0,
|
|
+ errors::InvalidArgument("input, filter_sizes and out_backprop tensors "
|
|
+ "must all have at least 1 element"));
|
|
+
|
|
const size_t shard_size =
|
|
(target_working_set_size + work_unit_size - 1) / work_unit_size;
|
|
|