From 9f0fef27dddc29efbd4b8b449f1f91ebef0dc000 Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EB=B0=95=EC=A2=85=ED=98=84/=EB=8F=99=EC=9E=91=EC=A0=9C?=
 =?utf8?q?=EC=96=B4Lab=28SR=29/Staff=20Engineer/=EC=82=BC=EC=84=B1?=
 =?utf8?q?=EC=A0=84=EC=9E=90?= <jh1302.park@samsung.com>
Date: Tue, 4 Sep 2018 16:41:11 +0900
Subject: [PATCH] [enco.caffe] Infer the output shape of Pooling with stride
 (#1322)

With this commit, PoolingSpec is now able to infer the output shape of
Pooling layer with stride values.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>
---
 contrib/enco/frontend/caffe/src/PoolingSpec.cpp    | 37 +++++++++++++++++-
 contrib/enco/frontend/caffe/src/PoolingSpec.h      |  4 ++
 .../enco/frontend/caffe/src/PoolingSpec.test.cpp   | 45 ++++++++++++++++++++++
 3 files changed, 84 insertions(+), 2 deletions(-)

diff --git a/contrib/enco/frontend/caffe/src/PoolingSpec.cpp b/contrib/enco/frontend/caffe/src/PoolingSpec.cpp
index e9b316d..7c5915e 100644
--- a/contrib/enco/frontend/caffe/src/PoolingSpec.cpp
+++ b/contrib/enco/frontend/caffe/src/PoolingSpec.cpp
@@ -37,6 +37,36 @@ uint32_t PoolingSpec::window_width(void) const
   return _param.kernel_size();
 }
 
+uint32_t PoolingSpec::vertical_stride(void) const
+{
+  if (_param.has_stride_h())
+  {
+    return _param.stride_h();
+  }
+
+  if (_param.has_stride())
+  {
+    return _param.stride();
+  }
+
+  return 1;
+}
+
+uint32_t PoolingSpec::horizontal_stride(void) const
+{
+  if (_param.has_stride_w())
+  {
+    return _param.stride_w();
+  }
+
+  if (_param.has_stride())
+  {
+    return _param.stride();
+  }
+
+  return 1;
+}
+
 nncc::core::ADT::tensor::Shape PoolingSpec::ofm_shape(void) const
 {
   nncc::core::ADT::tensor::Shape res;
@@ -51,7 +81,10 @@ nncc::core::ADT::tensor::Shape PoolingSpec::ofm_shape(void) const
   res.dim(1) = _ifm_shape.dim(1);
 
   // H and W are derived from IFM, Window, and Padding
-  res.dim(2) = _ifm_shape.dim(2) - window_height() + 1;
-  res.dim(3) = _ifm_shape.dim(3) - window_width() + 1;
+  // TODO Remove the following asserts
+  assert((_ifm_shape.dim(2) - window_height()) % vertical_stride() == 0);
+  assert((_ifm_shape.dim(3) - window_width()) % horizontal_stride() == 0);
+  res.dim(2) = (_ifm_shape.dim(2) - window_height()) / vertical_stride() + 1;
+  res.dim(3) = (_ifm_shape.dim(3) - window_width()) / horizontal_stride() + 1;
   return res;
 }
diff --git a/contrib/enco/frontend/caffe/src/PoolingSpec.h b/contrib/enco/frontend/caffe/src/PoolingSpec.h
index 8993501..d2ffa3a 100644
--- a/contrib/enco/frontend/caffe/src/PoolingSpec.h
+++ b/contrib/enco/frontend/caffe/src/PoolingSpec.h
@@ -19,6 +19,10 @@ public:
   uint32_t window_width(void) const;
 
 public:
+  uint32_t vertical_stride(void) const;
+  uint32_t horizontal_stride(void) const;
+
+public:
   nncc::core::ADT::tensor::Shape ofm_shape(void) const;
 
 private:
diff --git a/contrib/enco/frontend/caffe/src/PoolingSpec.test.cpp b/contrib/enco/frontend/caffe/src/PoolingSpec.test.cpp
index c36e9aa..0372896 100644
--- a/contrib/enco/frontend/caffe/src/PoolingSpec.test.cpp
+++ b/contrib/enco/frontend/caffe/src/PoolingSpec.test.cpp
@@ -147,3 +147,48 @@ TEST_F(PoolingSpecTest, kernel_size_same_for_all)
     ASSERT_EQ(expected, obtained);
   }
 }
+
+TEST_F(PoolingSpecTest, stride_for_all)
+{
+  const tensor::Shape ifm_shape{1, 3, 15, 15};
+
+  ::caffe::NetParameter param;
+  {
+    SequentialBuilder builder{&param};
+
+    builder.addInputLayer(ifm_shape);
+
+    // clang-format off
+    const char *prototxt = STRING(
+      name : "pool"
+      type : "Pooling"
+      bottom : "data"
+      top : "pool"
+      pooling_param {
+        pool: MAX
+        kernel_size : 3
+        stride: 2
+      }
+    );
+    // clang-format on
+
+    builder.addLayer(prototxt);
+  }
+
+  ::caffe::Net<float> net{param};
+
+  PoolingSpec spec{param.layer(1).pooling_param()};
+
+  spec.ifm_shape(ifm_shape);
+
+  ASSERT_EQ(spec.vertical_stride(), 2);
+  ASSERT_EQ(spec.horizontal_stride(), 2);
+
+  // Check 'ofm_shape'
+  {
+    auto expected = as_tensor_shape(net.blob_by_name("pool")->shape());
+    auto obtained = spec.ofm_shape();
+
+    ASSERT_EQ(expected, obtained);
+  }
+}
-- 
2.7.4