zzos/tensorflow
1
0
Fork 0
mirror of https://github.com/tensorflow/tensorflow.git synced 2026-09-28 05:13:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Handle zero-sized inputs in tf.nn.lrn and its gradient op to prevent GPU crash.

Browse Source 
PiperOrigin-RevId: 987255629
This commit is contained in:
Dmitri Latushko
2026-09-25 03:08:22 -07:00
committed by TensorFlower Gardener
parent ea22af286c
commit 17c17ea4f0
4 changed files with 109 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/core/kernels/lrn_op.cc
+12
View File
@@ -352,6 +352,12 @@ class LRNOp : public OpKernel {
context->allocate_output(
0, TensorShape({batch, rows, cols, depth}), &output));
// NOMUTANTS(LCR): Early return is a no-op guard on CPU but prevents fatal
// cuDNN crashes on GPU.
if (in.NumElements() == 0) {
return;
}
LaunchLRN<Device, T> launcher(depth_radius_, bias_, alpha_, beta_);
launcher.launch(context, this, in, output);
}
@@ -696,6 +702,12 @@ class LRNGradOp : public OpKernel {
context->allocate_output(
0, TensorShape({batch, rows, cols, depth}), &output));
// NOMUTANTS(LCR): Early return is a no-op guard on CPU but prevents fatal
// cuDNN crashes on GPU.
if (in_grads.NumElements() == 0) {
return;
}
LaunchLRNGrad<Device, T> launcher(depth_radius_, bias_, alpha_, beta_);
launcher.launch(context, this, in_grads, in_image, out_image, output);
}
tensorflow/core/kernels/lrn_op_test.cc
+38
View File
@@ -159,6 +159,44 @@ TEST_F(LRNFloatTest, Depth16) {
EXPECT_TRUE(Compare());
}
TEST_F(LRNFloatTest, ZeroSizeInput) {
TF_ASSERT_OK(NodeDefBuilder("lrn_op", "LRN")
.Input(FakeInput())
.Attr("depth_radius", 2)
.Attr("bias", 1.0f)
.Attr("alpha", 1.0f)
.Attr("beta", 0.5f)
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInput<float>(TensorShape({0, 0, 0, 0}),
[](int i) -> float { return 0.0f; });
TF_ASSERT_OK(RunOpKernel());
EXPECT_EQ(GetOutput(0)->NumElements(), 0);
EXPECT_EQ(GetOutput(0)->shape(), TensorShape({0, 0, 0, 0}));
}
TEST_F(LRNFloatTest, ZeroSizeInputGrad) {
TF_ASSERT_OK(NodeDefBuilder("lrn_grad_op", "LRNGrad")
.Input(FakeInput())
.Input(FakeInput())
.Input(FakeInput())
.Attr("depth_radius", 2)
.Attr("bias", 1.0f)
.Attr("alpha", 1.0f)
.Attr("beta", 0.5f)
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInput<float>(TensorShape({0, 0, 0, 0}),
[](int i) -> float { return 0.0f; });
AddInput<float>(TensorShape({0, 0, 0, 0}),
[](int i) -> float { return 0.0f; });
AddInput<float>(TensorShape({0, 0, 0, 0}),
[](int i) -> float { return 0.0f; });
TF_ASSERT_OK(RunOpKernel());
EXPECT_EQ(GetOutput(0)->NumElements(), 0);
EXPECT_EQ(GetOutput(0)->shape(), TensorShape({0, 0, 0, 0}));
}
static double RndGaussian(random::SimplePhilox* rnd) {
// Box-Muller transformation.
// See, for example, http://www.taygeta.com/random/gaussian.html
tensorflow/python/kernel_tests/nn_ops/BUILD
+1
View File
@@ -596,6 +596,7 @@ cuda_py_strict_test(
size = "medium",
srcs = ["lrn_op_test.py"],
deps = [
"//tensorflow/python/eager:backprop",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:for_generated_wrappers",
tensorflow/python/kernel_tests/nn_ops/lrn_op_test.py
+58
View File
@@ -18,6 +18,7 @@ import copy
import numpy as np
from tensorflow.python.eager import backprop
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
@@ -191,6 +192,63 @@ class LRNOpTest(test.TestCase):
if not test.is_gpu_available():
self._RunAndVerifyGradients(dtypes.float16)
@test_util.run_in_graph_and_eager_modes
def testZeroSizeInput(self) -> None:
for shape in [
[0, 0, 0, 0],
[2, 0, 4, 3],
[0, 3, 3, 2],
[1, 2, 0, 3],
[1, 2, 3, 0],
]:
x = array_ops.zeros(shape, dtype=dtypes.float32)
y = nn.local_response_normalization(
x, depth_radius=2, bias=1.0, alpha=1.0, beta=0.5
)
result = self.evaluate(y)
self.assertEqual(result.shape, tuple(shape))
@test_util.run_in_graph_and_eager_modes
def testZeroSizeGradInput(self) -> None:
for shape in [
[0, 0, 0, 0],
[2, 0, 4, 3],
[0, 3, 3, 2],
[1, 2, 0, 3],
[1, 2, 3, 0],
]:
x = array_ops.zeros(shape, dtype=dtypes.float32)
y = nn.lrn_grad(
input_grads=x,
input_image=x,
output_image=x,
depth_radius=2,
bias=1.0,
alpha=1.0,
beta=0.5,
)
result = self.evaluate(y)
self.assertEqual(result.shape, tuple(shape))
@test_util.run_in_graph_and_eager_modes
def testZeroSizeAutodiff(self) -> None:
for shape in [
[0, 0, 0, 0],
[2, 0, 4, 3],
[0, 3, 3, 2],
[1, 2, 0, 3],
[1, 2, 3, 0],
]:
x = array_ops.zeros(shape, dtype=dtypes.float32)
with backprop.GradientTape() as tape:
tape.watch(x)
y = nn.local_response_normalization(
x, depth_radius=2, bias=1.0, alpha=1.0, beta=0.5
)
grad = tape.gradient(y, x)
result = self.evaluate(grad)
self.assertEqual(result.shape, tuple(shape))
if __name__ == "__main__":
test.main()