52 lines
2.1 KiB
Diff
52 lines
2.1 KiB
Diff
|
From 5f7975d09eac0f10ed8a17dbb6f5964977725adc Mon Sep 17 00:00:00 2001
|
||
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
||
|
|
Date: Tue, 27 Apr 2021 17:45:43 -0700
|
||
|
|
Subject: [PATCH] Prevent another div by 0 in optimized pooling implementations
|
||
|
|
TFLite
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 370800091
|
||
|
|
Change-Id: I2119352f57fb5ca4f2051e0e2d749403304a979b
|
||
|
|
---
|
||
|
|
tensorflow/lite/kernels/pooling.cc | 4 ++++
|
||
|
|
tensorflow/lite/kernels/pooling_test.cc | 13 +++++++++++++
|
||
|
|
2 files changed, 17 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/lite/kernels/pooling.cc b/tensorflow/lite/kernels/pooling.cc
|
||
|
|
index 1ae3d207b135e..474bd3825f4ff 100644
|
||
|
|
--- a/tensorflow/lite/kernels/pooling.cc
|
||
|
|
+++ b/tensorflow/lite/kernels/pooling.cc
|
||
|
|
@@ -87,6 +87,10 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) {
|
||
|
|
auto padding = params->padding;
|
||
|
|
int out_width, out_height;
|
||
|
|
|
||
|
|
+ // Prevent division by 0 in optimized pooling implementations
|
||
|
|
+ TF_LITE_ENSURE(context, params->stride_height > 0);
|
||
|
|
+ TF_LITE_ENSURE(context, params->stride_width > 0);
|
||
|
|
+
|
||
|
|
data->padding = ComputePaddingHeightWidth(
|
||
|
|
params->stride_height, params->stride_width, 1, 1, height, width,
|
||
|
|
params->filter_height, params->filter_width, padding, &out_height,
|
||
|
|
diff --git a/tensorflow/lite/kernels/pooling_test.cc b/tensorflow/lite/kernels/pooling_test.cc
|
||
|
|
index e614fedccfd50..108195388141d 100644
|
||
|
|
--- a/tensorflow/lite/kernels/pooling_test.cc
|
||
|
|
+++ b/tensorflow/lite/kernels/pooling_test.cc
|
||
|
|
@@ -1151,5 +1151,18 @@ TEST(FloatPoolingOpTest, L2PoolPaddingValidSlide1) {
|
||
|
|
EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.5, 6.0, 6.5}));
|
||
|
|
}
|
||
|
|
|
||
|
|
+#ifdef GTEST_HAS_DEATH_TEST
|
||
|
|
+TEST(FloatPoolingOpTest, MaxPoolWithZeroStride) {
|
||
|
|
+ EXPECT_DEATH(
|
||
|
|
+ FloatPoolingOpModel m(BuiltinOperator_MAX_POOL_2D,
|
||
|
|
+ /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}},
|
||
|
|
+ /*filter_width=*/2, /*filter_height=*/2,
|
||
|
|
+ /*output=*/{TensorType_FLOAT32, {}},
|
||
|
|
+ /*padding=*/Padding_VALID,
|
||
|
|
+ /*stride_w=*/0, /*stride_h=*/0),
|
||
|
|
+ "Cannot allocate tensors");
|
||
|
|
+}
|
||
|
|
+#endif
|
||
|
|
+
|
||
|
|
} // namespace
|
||
|
|
} // namespace tflite
|