[layer] added start/end dimension in flatten layer

 - For now flatten layer flatten all dimension except batch.
   This commit will be able to flatten only the sub dimensions

Signed-off-by: hyeonseok <hs89.lee@samsung.com>

diff --git a/nntrainer/layers/common_properties.cpp b/nntrainer/layers/common_properties.cpp
index 5e6f6669ae26e1929872c83191770b81f8eb7a05..737d47609fb38329f60cf06f0800c4c8ba267406 100644 (file)
--- a/nntrainer/layers/common_properties.cpp
+++ b/nntrainer/layers/common_properties.cpp
@@ -100,6 +100,18 @@ bool Axis::isValid(const unsigned int &value) const {
   return value < ml::train::TensorDim::MAXDIM;
 }
 
+StartDimension::StartDimension(unsigned int value) { set(value); }
+
+bool StartDimension::isValid(const unsigned int &value) const {
+  return value > 0 && value < ml::train::TensorDim::MAXDIM;
+}
+
+EndDimension::EndDimension(unsigned int value) { set(value); }
+
+bool EndDimension::isValid(const unsigned int &value) const {
+  return value > 0 && value < ml::train::TensorDim::MAXDIM;
+}
+
 bool SplitDimension::isValid(const unsigned int &value) const {
   return value > 0 && value < ml::train::TensorDim::MAXDIM;
 }

diff --git a/nntrainer/layers/common_properties.h b/nntrainer/layers/common_properties.h
index 33802beacfd7358cee3343c40f3000282d468625..12368a38b1a9b88f3d7ade97d4c3c3f2c2b03fc6 100644 (file)
--- a/nntrainer/layers/common_properties.h
+++ b/nntrainer/layers/common_properties.h
@@ -307,6 +307,50 @@ public:
   bool isValid(const unsigned int &value) const override;
 };
 
+/**
+ * @brief StartDimension property, start dimension to be flatten
+ *
+ */
+class StartDimension : public Axis {
+public:
+  StartDimension(unsigned int value = 1);
+  static constexpr const char *key = "start_dimension";
+  using prop_tag = uint_prop_tag;
+
+  /**
+   * @brief check if given value is valid
+   *
+   * @param v value to check
+   * @retval true if it is greater than 0 and smaller than
+   * ml::train::TensorDim::MAXDIM
+   * @retval false if it is smaller or equal to 0 or greater than
+   * ml::train::TensorDim::MAXDIM
+   */
+  bool isValid(const unsigned int &value) const override;
+};
+
+/**
+ * @brief EndDimension property, end dimension to be flatten
+ *
+ */
+class EndDimension : public Axis {
+public:
+  EndDimension(unsigned int value = ml::train::TensorDim::MAXDIM - 1);
+  static constexpr const char *key = "end_dimension";
+  using prop_tag = uint_prop_tag;
+
+  /**
+   * @brief check if given value is valid
+   *
+   * @param v value to check
+   * @retval true if it is greater than 0 and smaller than
+   * ml::train::TensorDim::MAXDIM
+   * @retval false if it is smaller or equal to 0 or greater than
+   * ml::train::TensorDim::MAXDIM
+   */
+  bool isValid(const unsigned int &value) const override;
+};
+
 /**
  * @brief SplitDimension property, dimension along which to split the input
  *

diff --git a/nntrainer/layers/flatten_layer.cpp b/nntrainer/layers/flatten_layer.cpp
index 286182f00f7543e12f6b98b698766409f4c1e43b..413e32c061db99a486a68704c132b06b47c48554 100644 (file)
--- a/nntrainer/layers/flatten_layer.cpp
+++ b/nntrainer/layers/flatten_layer.cpp
@@ -6,6 +6,7 @@
  * @date   16 June 2020
  * @see    https://github.com/nnstreamer/nntrainer
  * @author Jijoong Moon <jijoong.moon@samsung.com>
+ * @author hyeonseok Lee <hs89.lee@samsung.com>
  * @bug           No known bugs except for NYI items
  * @brief  This is Flatten Layer Class for Neural Network
  *
@@ -25,8 +26,29 @@ static constexpr size_t SINGLE_INOUT_IDX = 0;
 void FlattenLayer::finalize(InitLayerContext &context) {
   const TensorDim &in_dim = context.getInputDimensions()[0];
 
-  std::string target_shape =
-    "target_shape=1:1:" + std::to_string(in_dim.getFeatureLen());
+  std::string target_shape;
+
+  const unsigned int start_dimension =
+    std::get<props::StartDimension>(flatten_props).get();
+  const unsigned int end_dimension =
+    std::get<props::EndDimension>(flatten_props).get();
+
+  NNTR_THROW_IF(start_dimension > end_dimension, std::invalid_argument)
+    << "start_dimension is bigger than end_dimension";
+
+  TensorDim target_dim = in_dim;
+
+  unsigned int flattened_size = 1;
+  for (unsigned int i = start_dimension; i <= end_dimension; ++i) {
+    flattened_size *= in_dim[i];
+    target_dim[i] = 1;
+  }
+  target_dim[end_dimension] = flattened_size;
+
+  target_shape = "target_shape=" + std::to_string(target_dim[1]) + ":" +
+                 std::to_string(target_dim[2]) + ":" +
+                 std::to_string(target_dim[3]);
+
   ReshapeLayer::setProperty({target_shape});
 
   /** @note the output dimension is in invalid state till finalize of
@@ -39,7 +61,8 @@ void FlattenLayer::finalize(InitLayerContext &context) {
 }
 
 void FlattenLayer::setProperty(const std::vector<std::string> &values) {
-  auto remain_props = loadProperties(values, reshape_props);
+  auto remain_props = loadProperties(values, flatten_props);
+  remain_props = loadProperties(remain_props, reshape_props);
   if (!remain_props.empty()) {
     std::string msg = "[FlattenLayer] Unknown Layer Properties count " +
                       std::to_string(values.size());

diff --git a/nntrainer/layers/flatten_layer.h b/nntrainer/layers/flatten_layer.h
index 66ea6553b66bd26bdd8242dbe075318276555dcf..238f67140edb3ab7f88576b82afd024ee82394f5 100644 (file)
--- a/nntrainer/layers/flatten_layer.h
+++ b/nntrainer/layers/flatten_layer.h
@@ -28,7 +28,8 @@ public:
   /**
    * @brief     Constructor of Flatten Layer
    */
-  FlattenLayer() : ReshapeLayer() {}
+  FlattenLayer() : ReshapeLayer(), flatten_props(
+    props::StartDimension(), props::EndDimension()) {}
 
   /**
    * @brief     Destructor of Flatten Layer
@@ -70,6 +71,8 @@ public:
   const std::string getType() const override { return FlattenLayer::type; };
 
   inline static const std::string type = "flatten";
+
+  std::tuple<props::StartDimension, props::EndDimension> flatten_props;
 };
 
 } // namespace nntrainer

diff --git a/test/unittest/layers/unittest_layers_flatten.cpp b/test/unittest/layers/unittest_layers_flatten.cpp
index da55901f59d46f388f9f67a2728a10504d0a5893..30c147d6a7f197416d89ef043135906bb825e4ee 100644 (file)
--- a/test/unittest/layers/unittest_layers_flatten.cpp
+++ b/test/unittest/layers/unittest_layers_flatten.cpp
@@ -21,5 +21,17 @@ auto semantic_flatten = LayerSemanticsParamType(
   nntrainer::FlattenLayer::type, {},
   LayerCreateSetPropertyOptions::AVAILABLE_FROM_APP_CONTEXT, false, 1);
 
-GTEST_PARAMETER_TEST(Flatten, LayerSemantics,
-                     ::testing::Values(semantic_flatten));
+auto semantic_flatten_with_start_dim = LayerSemanticsParamType(
+  nntrainer::createLayer<nntrainer::FlattenLayer>,
+  nntrainer::FlattenLayer::type, {"start_dimension = 2"},
+  LayerCreateSetPropertyOptions::AVAILABLE_FROM_APP_CONTEXT, false, 1);
+
+auto semantic_flatten_with_start_dim_end_dim = LayerSemanticsParamType(
+  nntrainer::createLayer<nntrainer::FlattenLayer>,
+  nntrainer::FlattenLayer::type, {"start_dimension = 2", "end_dimension = 3"},
+  LayerCreateSetPropertyOptions::AVAILABLE_FROM_APP_CONTEXT, false, 1);
+
+GTEST_PARAMETER_TEST(
+  Flatten, LayerSemantics,
+  ::testing::Values(semantic_flatten, semantic_flatten_with_start_dim,
+                    semantic_flatten_with_start_dim_end_dim));