From 70f3a3afeaecbb1c6ea4b99b96a47fdf888110eb Mon Sep 17 00:00:00 2001
From: Parichay Kapoor <pk.kapoor@samsung.com>
Date: Fri, 22 Jan 2021 12:41:14 +0900
Subject: [PATCH] [optimization] Bug fix for in-place layer optimization

Inplace layer optimization is performed for multiple layers - activation and batch normalization layers
and this list will increase with data augmentation etc.
However, the in-place layers cannot work correctly consecutively if these layers are trainable.
They can work perfectly is they dont need to pass the derivative back.

For now, this patch limits two consecutive layers to be in-place.
This will be made generic later dependent on the trainable and inPlace property of the layer.

**Self evaluation:**
1. Build test: [x]Passed [ ]Failed [ ]Skipped
2. Run test: [x]Passed [ ]Failed [ ]Skipped

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>
---
 nntrainer/graph/network_graph.cpp           | 28 ++++++++++++++++++++++++----
 test/unittest/unittest_nntrainer_models.cpp |  8 +++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/nntrainer/graph/network_graph.cpp b/nntrainer/graph/network_graph.cpp
index 4ffd520..0503028 100644
--- a/nntrainer/graph/network_graph.cpp
+++ b/nntrainer/graph/network_graph.cpp
@@ -28,6 +28,13 @@
 
 namespace nntrainer {
 
+/**
+ * @todo Make inPlace as a static property of the layer and a state to verify if
+ * this layer is working in-place
+ */
+static const std::vector<std::string> in_place_layers = {
+  ActivationLayer::type, BatchNormalizationLayer::type};
+
 void NetworkGraph::updateNameInLayers(
   std::vector<std::shared_ptr<Layer>> &layers, const std::string &cname,
   const std::string &name) {
@@ -552,9 +559,22 @@ void NetworkGraph::inPlaceOptimize(const std::string &layer_type,
       if (loc == prev_layer->output_layers.size())
         throw std::runtime_error("Internal error in the formed graph.");
 
-      if (prev_layer->getType() == InputLayer::type ||
-          prev_layer->getType() == ActivationLayer::type)
+      if (prev_layer->getType() == InputLayer::type)
+        continue;
+
+      /** check if previous layer was also in-place */
+      bool prev_layer_in_place = false;
+      for (auto const &in_place_layer : in_place_layers) {
+        if (prev_layer->getType() == in_place_layer) {
+          prev_layer_in_place = true;
+          break;
+        }
+      }
+
+      /** Two layers cant work in-place consecutively */
+      if (prev_layer_in_place)
         continue;
+
       /** Share tensor with next layer */
       l->net_input[0] = l->net_hidden[0];
       prev_layer->net_hidden[loc] = l->net_hidden[0];
@@ -566,8 +586,8 @@ void NetworkGraph::inPlaceOptimize(const std::string &layer_type,
 }
 
 void NetworkGraph::inPlaceOptimize(Manager &manager) {
-  inPlaceOptimize(BatchNormalizationLayer::type, manager);
-  inPlaceOptimize(ActivationLayer::type, manager);
+  for (auto const &layer_type : in_place_layers)
+    inPlaceOptimize(layer_type, manager);
 }
 
 } /* namespace nntrainer */
diff --git a/test/unittest/unittest_nntrainer_models.cpp b/test/unittest/unittest_nntrainer_models.cpp
index 6f3270f..de76dcd 100644
--- a/test/unittest/unittest_nntrainer_models.cpp
+++ b/test/unittest/unittest_nntrainer_models.cpp
@@ -247,7 +247,13 @@ void NodeWatcher::forward(int iteration, NodeWatcher &next_node) {
 
   std::vector<nntrainer::Tensor> out = node.layer->getOutputs();
 
-  if (next_node.node.layer->getType() != nntrainer::ActivationLayer::type)
+  /**
+   * @todo Do not veify if the layer is operting in-place by checking its
+   * property
+   */
+  if (next_node.node.layer->getType() != nntrainer::ActivationLayer::type &&
+      next_node.node.layer->getType() !=
+        nntrainer::BatchNormalizationLayer::type)
     verify(out[0], expected_output, err_msg + " at output");
 }
 
-- 
2.7.4