From: Parichay Kapoor <pk.kapoor@samsung.com>
Date: Thu, 8 Apr 2021 10:30:22 +0000 (+0900)
Subject: [graph] Update graph node
X-Git-Tag: submit/tizen/20210827.122527~382
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=a63273f39165f0350cbae4cd86ff2464905db7ca;p=platform%2Fcore%2Fml%2Fnntrainer.git

[graph] Update graph node

Update graph node to use the updated LayerNode
Also add minor bug fix for the graph

**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>
---

diff --git a/Applications/Custom/LayerClient/jni/main.cpp b/Applications/Custom/LayerClient/jni/main.cpp
index 5e1cc984..4352e16e 100644
--- a/Applications/Custom/LayerClient/jni/main.cpp
+++ b/Applications/Custom/LayerClient/jni/main.cpp
@@ -199,7 +199,7 @@ int main(int argc, char *argv[]) {
     /// registerFactory excepts a function that returns unique_ptr<Layer> from
     /// std::vector<std::string> ml::train::createLayer<T> is a templated
     /// function for generic usage
-    app_context.registerFactory(ml::train::createLayer<custom::PowLayer>);
+    app_context.registerFactory(nntrainer::createLayer<custom::PowLayer>);
   } catch (std::invalid_argument &e) {
     std::cerr << "failed to register factory, reason: " << e.what()
               << std::endl;
diff --git a/Applications/Custom/LayerPlugin/layer_plugin_test.cpp b/Applications/Custom/LayerPlugin/layer_plugin_test.cpp
index 37806fa2..0b21b3ba 100644
--- a/Applications/Custom/LayerPlugin/layer_plugin_test.cpp
+++ b/Applications/Custom/LayerPlugin/layer_plugin_test.cpp
@@ -19,7 +19,7 @@
 
 #include <app_context.h>
 #include <layer.h>
-#include <layer_internal.h>
+#include <layer_node.h>
 
 const char *NNTRAINER_PATH = std::getenv("NNTRAINER_PATH");
 
@@ -30,7 +30,7 @@ TEST(AppContext, DlRegisterOpen_p) {
 
   ac.registerLayer("libpow_layer.so", NNTRAINER_PATH);
 
-  auto layer = ac.createObject<ml::train::Layer>("pow");
+  auto layer = ac.createObject<nntrainer::Layer>("pow");
 
   EXPECT_EQ(layer->getType(), "pow");
 }
@@ -50,7 +50,7 @@ TEST(AppContext, DlRegisterDirectory_p) {
 
   ac.registerLayerFromDirectory(NNTRAINER_PATH);
 
-  auto layer = ac.createObject<ml::train::Layer>("pow");
+  auto layer = ac.createObject<nntrainer::Layer>("pow");
 
   EXPECT_EQ(layer->getType(), "pow");
 }
@@ -65,11 +65,10 @@ TEST(AppContext, DlRegisterDirectory_n) {
 TEST(AppContext, DefaultEnvironmentPath_p) {
   /// as NNTRAINER_PATH is fed to the test, this should success without an
   /// error
-  auto l = ml::train::createLayer("pow");
+  std::shared_ptr<ml::train::Layer> l = ml::train::createLayer("pow");
   EXPECT_EQ(l->getType(), "pow");
 
-  std::unique_ptr<nntrainer::Layer> layer(
-    static_cast<nntrainer::Layer *>(l.release()));
+  auto layer = nntrainer::getLayerDevel(l);
 
   std::ifstream input_file("does_not_exist");
   EXPECT_NO_THROW(layer->read(input_file));
diff --git a/Applications/Custom/pow.cpp b/Applications/Custom/pow.cpp
index 28d84752..185e0662 100644
--- a/Applications/Custom/pow.cpp
+++ b/Applications/Custom/pow.cpp
@@ -21,6 +21,9 @@ namespace custom {
 const std::string PowLayer::type = "pow";
 namespace PowUtil {
 
+/**
+ * @brief Entry structure for handling properties
+ */
 struct Entry {
   std::string key;
   std::string value;
@@ -136,13 +139,13 @@ void PowLayer::calcDerivative() {
 
 #ifdef PLUGGABLE
 
-ml::train::Layer *create_pow_layer() {
+nntrainer::Layer *create_pow_layer() {
   auto layer = new PowLayer();
   std::cout << "power created\n";
   return layer;
 }
 
-void destory_pow_layer(ml::train::Layer *layer) {
+void destory_pow_layer(nntrainer::Layer *layer) {
   std::cout << "power deleted\n";
   delete layer;
 }
diff --git a/Applications/Embedding/jni/main.cpp b/Applications/Embedding/jni/main.cpp
index 1322b259..c1e12e59 100644
--- a/Applications/Embedding/jni/main.cpp
+++ b/Applications/Embedding/jni/main.cpp
@@ -15,16 +15,18 @@
  */
 
 #include <cmath>
-#include <dataset.h>
 #include <fstream>
 #include <iostream>
-#include <ml-api-common.h>
-#include <neuralnet.h>
 #include <sstream>
 #include <stdlib.h>
 #include <tensor.h>
 #include <time.h>
 
+#include <dataset.h>
+#include <layer_node.h>
+#include <ml-api-common.h>
+#include <neuralnet.h>
+
 std::string data_file;
 
 const unsigned int total_train_data_size = 10;
@@ -181,8 +183,10 @@ int main(int argc, char *argv[]) {
    */
   std::vector<std::vector<float>> inputVector, outputVector;
   nntrainer::NeuralNetwork NN;
-  std::shared_ptr<nntrainer::Layer> layer;
-  std::shared_ptr<nntrainer::Layer> layer_fc;
+  std::shared_ptr<ml::train::Layer> layer;
+  std::shared_ptr<ml::train::Layer> layer_fc;
+  std::shared_ptr<nntrainer::Layer> layer_;
+  std::shared_ptr<nntrainer::Layer> layer_fc_;
   std::string name = "embedding";
   std::string fc_name = "outputlayer";
   nntrainer::Tensor weight;
@@ -205,7 +209,10 @@ int main(int argc, char *argv[]) {
     NN.getLayer(name.c_str(), &layer);
     NN.getLayer(fc_name.c_str(), &layer_fc);
 
-    weight = layer->getWeights()[0].getVariable();
+    layer_ = nntrainer::getLayerDevel(layer);
+    layer_fc_ = nntrainer::getLayerDevel(layer_fc);
+
+    weight = layer_->getWeights()[0].getVariable();
     weight.print(std::cout);
 
   } catch (...) {
@@ -226,11 +233,11 @@ int main(int argc, char *argv[]) {
     // std::string name = "embedding";
     // NN.getLayer(name.c_str(), &layer);
 
-    weight = layer->getWeights()[0].getVariable();
+    weight = layer_->getWeights()[0].getVariable();
     weight.print(std::cout);
 
     nntrainer::Tensor golden(1, 1, 15, 8);
-    
+
     try {
       loadFile("embedding_weight_golden.out", golden);
       golden.print(std::cout);
@@ -238,7 +245,8 @@ int main(int argc, char *argv[]) {
       nntrainer::Tensor weight_out_fc(1, 1, 32, 1);
       loadFile("fc_weight_golden.out", weight_out_fc);
       weight_out_fc.print(std::cout);
-      weight_fc = layer_fc->getWeights()[0].getVariable();
+
+      weight_fc = layer_fc_->getWeights()[0].getVariable();
       weight_fc.print(std::cout);
     } catch (...) {
       std::cerr << "Error during loadFile\n";
diff --git a/Applications/ReinforcementLearning/DeepQ/jni/main.cpp b/Applications/ReinforcementLearning/DeepQ/jni/main.cpp
index ee1fe38e..f51dceb4 100644
--- a/Applications/ReinforcementLearning/DeepQ/jni/main.cpp
+++ b/Applications/ReinforcementLearning/DeepQ/jni/main.cpp
@@ -514,7 +514,7 @@ int main(int argc, char **argv) {
         writeFile << "mainNet Loss : " << mainNet.getLoss()
                   << " : targetNet Loss : " << targetNet.getLoss() << "\n";
         std::cout << "\n\n =================== TRAINIG & COPY NET "
-                    "==================\n\n";
+                     "==================\n\n";
         std::cout << "mainNet Loss : ";
         std::cout.width(15);
         std::cout << mainNet.getLoss() << "\n targetNet Loss : ";
diff --git a/Applications/SimpleShot/task_runner.cpp b/Applications/SimpleShot/task_runner.cpp
index c6c8214e..6b38e878 100644
--- a/Applications/SimpleShot/task_runner.cpp
+++ b/Applications/SimpleShot/task_runner.cpp
@@ -223,11 +223,11 @@ int main(int argc, char **argv) {
 
   try {
     app_context.registerFactory(
-      ml::train::createLayer<simpleshot::layers::CenteringLayer>);
+      nntrainer::createLayer<simpleshot::layers::CenteringLayer>);
     app_context.registerFactory(
-      ml::train::createLayer<simpleshot::layers::L2NormLayer>);
+      nntrainer::createLayer<simpleshot::layers::L2NormLayer>);
     app_context.registerFactory(
-      ml::train::createLayer<simpleshot::layers::CentroidKNN>);
+      nntrainer::createLayer<simpleshot::layers::CentroidKNN>);
   } catch (std::exception &e) {
     std::cerr << "registering factory failed: " << e.what();
     return 1;
diff --git a/Applications/SimpleShot/test/simpleshot_centering_test.cpp b/Applications/SimpleShot/test/simpleshot_centering_test.cpp
index 921ece8a..2205f255 100644
--- a/Applications/SimpleShot/test/simpleshot_centering_test.cpp
+++ b/Applications/SimpleShot/test/simpleshot_centering_test.cpp
@@ -16,6 +16,7 @@
 #include <memory>
 
 #include <app_context.h>
+#include <layer_node.h>
 #include <manager.h>
 #include <nntrainer_test_util.h>
 
@@ -34,12 +35,12 @@ TEST(centering, simple_functions) {
   file.close();
 
   auto &app_context = nntrainer::AppContext::Global();
-  app_context.registerFactory(ml::train::createLayer<CenteringLayer>);
+  app_context.registerFactory(nntrainer::createLayer<CenteringLayer>);
 
-  auto c = app_context.createObject<ml::train::Layer>(
+  auto c = app_context.createObject<nntrainer::Layer>(
     "centering", {"feature_path=feature.bin", "input_shape=1:1:4"});
 
-  std::unique_ptr<CenteringLayer> layer(
+  std::shared_ptr<CenteringLayer> layer(
     static_cast<CenteringLayer *>(c.release()));
 
   nntrainer::Manager manager;
diff --git a/Applications/SimpleShot/test/simpleshot_centroid_knn.cpp b/Applications/SimpleShot/test/simpleshot_centroid_knn.cpp
index 4f1bb6f7..c7bc890e 100644
--- a/Applications/SimpleShot/test/simpleshot_centroid_knn.cpp
+++ b/Applications/SimpleShot/test/simpleshot_centroid_knn.cpp
@@ -25,9 +25,9 @@ namespace layers {
 
 TEST(centroid_knn, simple_functions) {
   auto &app_context = nntrainer::AppContext::Global();
-  app_context.registerFactory(ml::train::createLayer<CentroidKNN>);
+  app_context.registerFactory(nntrainer::createLayer<CentroidKNN>);
 
-  auto c = app_context.createObject<ml::train::Layer>(
+  auto c = app_context.createObject<nntrainer::Layer>(
     "centroid_knn", {"num_class=5", "input_shape=1:1:3"});
 
   std::unique_ptr<CentroidKNN> layer(static_cast<CentroidKNN *>(c.release()));
diff --git a/Applications/SimpleShot/test/simpleshot_l2norm_test.cpp b/Applications/SimpleShot/test/simpleshot_l2norm_test.cpp
index ef6cd031..6d42c5ec 100644
--- a/Applications/SimpleShot/test/simpleshot_l2norm_test.cpp
+++ b/Applications/SimpleShot/test/simpleshot_l2norm_test.cpp
@@ -15,6 +15,7 @@
 #include <memory>
 
 #include <app_context.h>
+#include <layer_node.h>
 #include <manager.h>
 #include <nntrainer_test_util.h>
 
@@ -25,13 +26,12 @@ namespace layers {
 
 TEST(l2norm, simple_functions) {
   auto &app_context = nntrainer::AppContext::Global();
-  app_context.registerFactory(ml::train::createLayer<L2NormLayer>);
+  app_context.registerFactory(nntrainer::createLayer<L2NormLayer>);
 
   auto c =
-    app_context.createObject<ml::train::Layer>("l2norm", {"input_shape=1:1:4"});
+    app_context.createObject<nntrainer::Layer>("l2norm", {"input_shape=1:1:4"});
 
-  std::unique_ptr<nntrainer::Layer> layer(
-    static_cast<nntrainer::Layer *>(c.release()));
+  std::unique_ptr<L2NormLayer> layer(static_cast<L2NormLayer *>(c.release()));
 
   nntrainer::Manager manager;
   manager.setInferenceInOutMemoryOptimization(false);
diff --git a/api/ccapi/src/factory.cpp b/api/ccapi/src/factory.cpp
index bc65269e..accfd982 100644
--- a/api/ccapi/src/factory.cpp
+++ b/api/ccapi/src/factory.cpp
@@ -41,7 +41,12 @@ std::unique_ptr<Layer> createLayer(const LayerType &type,
 std::unique_ptr<Layer> createLayer(const std::string &type,
                                    const std::vector<std::string> &properties) {
   auto &ac = nntrainer::AppContext::Global();
-  return ac.createObject<Layer>(type, properties);
+  std::shared_ptr<nntrainer::Layer> nntr_layer =
+    ac.createObject<nntrainer::Layer>(type, properties);
+  std::unique_ptr<nntrainer::LayerNode> layer =
+    std::make_unique<nntrainer::LayerNode>(nntr_layer);
+
+  return layer;
 }
 
 std::unique_ptr<Optimizer>
@@ -58,7 +63,9 @@ createOptimizer(const OptimizerType &type,
 static std::unique_ptr<Layer>
 createLoss(nntrainer::LossType type,
            const std::vector<std::string> &properties) {
-  std::unique_ptr<Layer> layer = nntrainer::createLoss(type);
+  std::shared_ptr<nntrainer::Layer> nntr_layer = nntrainer::createLoss(type);
+  std::unique_ptr<nntrainer::LayerNode> layer =
+    std::make_unique<nntrainer::LayerNode>(nntr_layer);
 
   if (layer->setProperty(properties) != ML_ERROR_NONE)
     throw std::invalid_argument("Set properties failed for layer");
diff --git a/jni/Android.mk b/jni/Android.mk
index dbe1b1ef..0ff00a31 100644
--- a/jni/Android.mk
+++ b/jni/Android.mk
@@ -102,6 +102,7 @@ NNTRAINER_SRCS := $(NNTRAINER_ROOT)/nntrainer/models/neuralnet.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/tensor/tensor_dim.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/tensor/blas_interface.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/layers/layer.cpp \
+                  $(NNTRAINER_ROOT)/nntrainer/layers/layer_node.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/layers/layer_factory.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/layers/input_layer.cpp \
                   $(NNTRAINER_ROOT)/nntrainer/layers/output_layer.cpp \
diff --git a/nntrainer/app_context.cpp b/nntrainer/app_context.cpp
index a3b02b02..47b2ea60 100644
--- a/nntrainer/app_context.cpp
+++ b/nntrainer/app_context.cpp
@@ -197,56 +197,56 @@ static void add_default_object(AppContext &ac) {
                      "unknown", OptType::UNKNOWN);
 
   using LayerType = ml::train::LayerType;
-  ac.registerFactory(ml::train::createLayer<InputLayer>, InputLayer::type,
+  ac.registerFactory(nntrainer::createLayer<InputLayer>, InputLayer::type,
                      LayerType::LAYER_IN);
-  ac.registerFactory(ml::train::createLayer<FullyConnectedLayer>,
+  ac.registerFactory(nntrainer::createLayer<FullyConnectedLayer>,
                      FullyConnectedLayer::type, LayerType::LAYER_FC);
-  ac.registerFactory(ml::train::createLayer<BatchNormalizationLayer>,
+  ac.registerFactory(nntrainer::createLayer<BatchNormalizationLayer>,
                      BatchNormalizationLayer::type, LayerType::LAYER_BN);
-  ac.registerFactory(ml::train::createLayer<Conv2DLayer>, Conv2DLayer::type,
+  ac.registerFactory(nntrainer::createLayer<Conv2DLayer>, Conv2DLayer::type,
                      LayerType::LAYER_CONV2D);
-  ac.registerFactory(ml::train::createLayer<Pooling2DLayer>,
+  ac.registerFactory(nntrainer::createLayer<Pooling2DLayer>,
                      Pooling2DLayer::type, LayerType::LAYER_POOLING2D);
-  ac.registerFactory(ml::train::createLayer<FlattenLayer>, FlattenLayer::type,
+  ac.registerFactory(nntrainer::createLayer<FlattenLayer>, FlattenLayer::type,
                      LayerType::LAYER_FLATTEN);
-  ac.registerFactory(ml::train::createLayer<ActivationLayer>,
+  ac.registerFactory(nntrainer::createLayer<ActivationLayer>,
                      ActivationLayer::type, LayerType::LAYER_ACTIVATION);
-  ac.registerFactory(ml::train::createLayer<AdditionLayer>, AdditionLayer::type,
+  ac.registerFactory(nntrainer::createLayer<AdditionLayer>, AdditionLayer::type,
                      LayerType::LAYER_ADDITION);
-  ac.registerFactory(ml::train::createLayer<OutputLayer>, OutputLayer::type,
+  ac.registerFactory(nntrainer::createLayer<OutputLayer>, OutputLayer::type,
                      LayerType::LAYER_MULTIOUT);
-  ac.registerFactory(ml::train::createLayer<ConcatLayer>, ConcatLayer::type,
+  ac.registerFactory(nntrainer::createLayer<ConcatLayer>, ConcatLayer::type,
                      LayerType::LAYER_CONCAT);
-  ac.registerFactory(ml::train::createLayer<LossLayer>, LossLayer::type,
+  ac.registerFactory(nntrainer::createLayer<LossLayer>, LossLayer::type,
                      LayerType::LAYER_LOSS);
-  ac.registerFactory(ml::train::createLayer<PreprocessFlipLayer>,
+  ac.registerFactory(nntrainer::createLayer<PreprocessFlipLayer>,
                      PreprocessFlipLayer::type,
                      LayerType::LAYER_PREPROCESS_FLIP);
-  ac.registerFactory(ml::train::createLayer<PreprocessTranslateLayer>,
+  ac.registerFactory(nntrainer::createLayer<PreprocessTranslateLayer>,
                      PreprocessTranslateLayer::type,
                      LayerType::LAYER_PREPROCESS_TRANSLATE);
 #ifdef ENABLE_NNSTREAMER_BACKBONE
-  ac.registerFactory(ml::train::createLayer<NNStreamerLayer>,
+  ac.registerFactory(nntrainer::createLayer<NNStreamerLayer>,
                      NNStreamerLayer::type,
                      LayerType::LAYER_BACKBONE_NNSTREAMER);
 #endif
 #ifdef ENABLE_TFLITE_BACKBONE
-  ac.registerFactory(ml::train::createLayer<TfLiteLayer>, TfLiteLayer::type,
+  ac.registerFactory(nntrainer::createLayer<TfLiteLayer>, TfLiteLayer::type,
                      LayerType::LAYER_BACKBONE_TFLITE);
 #endif
-  ac.registerFactory(ml::train::createLayer<EmbeddingLayer>,
+  ac.registerFactory(nntrainer::createLayer<EmbeddingLayer>,
                      EmbeddingLayer::type, LayerType::LAYER_EMBEDDING);
 
-  ac.registerFactory(ml::train::createLayer<RNNLayer>, RNNLayer::type,
+  ac.registerFactory(nntrainer::createLayer<RNNLayer>, RNNLayer::type,
                      LayerType::LAYER_RNN);
 
-  ac.registerFactory(ml::train::createLayer<LSTMLayer>, LSTMLayer::type,
+  ac.registerFactory(nntrainer::createLayer<LSTMLayer>, LSTMLayer::type,
                      LayerType::LAYER_LSTM);
 
-  ac.registerFactory(ml::train::createLayer<TimeDistLayer>, TimeDistLayer::type,
+  ac.registerFactory(nntrainer::createLayer<TimeDistLayer>, TimeDistLayer::type,
                      LayerType::LAYER_TIME_DIST);
 
-  ac.registerFactory(AppContext::unknownFactory<ml::train::Layer>, "unknown",
+  ac.registerFactory(AppContext::unknownFactory<nntrainer::Layer>, "unknown",
                      LayerType::LAYER_UNKNOWN);
 }
 
@@ -338,15 +338,15 @@ int AppContext::registerLayer(const std::string &library_path,
     << func_tag << "custom layer must specify type name, but it is empty";
   pluggable->destroyfunc(layer);
 
-  FactoryType<ml::train::Layer> factory_func =
+  FactoryType<nntrainer::Layer> factory_func =
     [pluggable](const PropsType &prop) {
-      std::unique_ptr<ml::train::Layer> layer =
+      std::unique_ptr<nntrainer::Layer> layer =
         std::make_unique<internal::PluggedLayer>(pluggable);
 
       return layer;
     };
 
-  return registerFactory<ml::train::Layer>(factory_func, type);
+  return registerFactory<nntrainer::Layer>(factory_func, type);
 }
 
 std::vector<int>
diff --git a/nntrainer/app_context.h b/nntrainer/app_context.h
index 4aff0bf4..192c9376 100644
--- a/nntrainer/app_context.h
+++ b/nntrainer/app_context.h
@@ -65,7 +65,10 @@ public:
 
   template <typename... Ts> using FactoryMap = std::tuple<IndexType<Ts>...>;
 
-  AppContext(){};
+  /**
+   * @brief   Default constructor
+   */
+  AppContext() = default;
 
   /**
    *
@@ -133,6 +136,19 @@ public:
    */
   const std::string getWorkingPath(const std::string &path = "");
 
+  /**
+   * @brief Factory register function, use this function to register custom
+   * object
+   *
+   * @tparam T object to create. Currently Optimizer, Layer is supported
+   * @param factory factory function that creates std::unique_ptr<T>
+   * @param key key to access the factory, if key is empty, try to find key by
+   * calling factory({})->getType();
+   * @param int_key key to access the factory by integer, if it is -1(default),
+   * the function automatically unsigned the key and return
+   * @return const int unique integer value to access the current factory
+   * @throw invalid argument when key and/or int_key is already taken
+   */
   template <typename T>
   const int registerFactory(const PtrFactoryType<T> factory,
                             const std::string &key = "",
@@ -262,7 +278,7 @@ public:
   }
 
 private:
-  FactoryMap<ml::train::Optimizer, ml::train::Layer> factory_map;
+  FactoryMap<ml::train::Optimizer, nntrainer::Layer> factory_map;
   std::string working_path_base;
 };
 
diff --git a/nntrainer/compiler/ini_interpreter.cpp b/nntrainer/compiler/ini_interpreter.cpp
index 2449805c..c5b1a13c 100644
--- a/nntrainer/compiler/ini_interpreter.cpp
+++ b/nntrainer/compiler/ini_interpreter.cpp
@@ -114,14 +114,14 @@ std::vector<std::string> section2properties(dictionary *ini,
  * @return std::shared_ptr<Layer> return layer object
  */
 template <typename T>
-std::shared_ptr<Layer>
+std::shared_ptr<LayerNode>
 section2layer(dictionary *ini, const std::string &sec_name,
               const AppContext &ac, const std::string &backbone_file = "") {
   throw std::invalid_argument("supported only with a tag for now");
 }
 
 template <>
-std::shared_ptr<Layer>
+std::shared_ptr<LayerNode>
 section2layer<PlainLayer>(dictionary *ini, const std::string &sec_name,
                           const AppContext &ac,
                           const std::string &backbone_file) {
@@ -132,24 +132,26 @@ section2layer<PlainLayer>(dictionary *ini, const std::string &sec_name,
     << FUNC_TAG << "section type is invalid for section name: " << sec_name;
 
   auto properties = section2properties(ini, sec_name);
-  std::shared_ptr<ml::train::Layer> layer_ =
-    ac.createObject<ml::train::Layer>(layer_type, properties);
+  std::shared_ptr<Layer> nntr_layer =
+    ac.createObject<Layer>(layer_type, properties);
 
-  auto layer = std::static_pointer_cast<Layer>(layer_);
+  auto layer = std::make_unique<LayerNode>(nntr_layer);
 
-  if (layer->getDistribute()) {
+  if (nntr_layer->getDistribute()) {
     ml_logd("This %s layer is going to distributed", sec_name.c_str());
     std::shared_ptr<Layer> dist_layer =
       nntrainer::createLayer(TimeDistLayer::type);
-    std::dynamic_pointer_cast<TimeDistLayer>(dist_layer)->setDistLayer(layer);
-    layer = dist_layer;
+    std::dynamic_pointer_cast<TimeDistLayer>(dist_layer)
+      ->setDistLayer(nntr_layer);
+
+    layer = std::make_unique<LayerNode>(dist_layer);
   }
 
   return layer;
 }
 
 template <>
-std::shared_ptr<Layer>
+std::shared_ptr<LayerNode>
 section2layer<BackboneLayer>(dictionary *ini, const std::string &sec_name,
                              const AppContext &ac,
                              const std::string &backbone_file) {
@@ -173,10 +175,9 @@ section2layer<BackboneLayer>(dictionary *ini, const std::string &sec_name,
 
   auto properties = section2properties(ini, sec_name);
   properties.push_back("modelfile=" + backbone_file);
-  std::shared_ptr<ml::train::Layer> layer_ =
-    ac.createObject<ml::train::Layer>(type, properties);
+  std::shared_ptr<Layer> nntr_layer = ac.createObject<Layer>(type, properties);
 
-  auto layer = std::static_pointer_cast<Layer>(layer_);
+  auto layer = std::make_unique<LayerNode>(nntr_layer);
 
   return layer;
 }
@@ -201,7 +202,7 @@ static bool graphSupported(const std::string &backbone_name) {
  * @param sec_name section name
  * @return std::vector<std::shared_ptr<Layer>> mergeable graph
  */
-std::vector<std::shared_ptr<Layer>>
+std::vector<std::shared_ptr<LayerNode>>
 getMergeableGraph(std::shared_ptr<const GraphRepresentation> graph,
                   dictionary *ini, const std::string &sec_name) {
   std::string input_layer =
@@ -229,10 +230,10 @@ getMergeableGraph(std::shared_ptr<const GraphRepresentation> graph,
     << scale_size;
 
   for (auto &layer : g) {
-    layer->setTrainable(trainable);
-    layer->resetDimension();
+    layer->getObject()->setTrainable(trainable);
+    layer->getObject()->resetDimension();
     if (scale_size != 1) {
-      layer->scaleSize(scale_size);
+      layer->getObject()->scaleSize(scale_size);
     }
     /** TODO #361: this needs update in model file to be of dictionary format */
     // if (preload) {
@@ -251,13 +252,15 @@ getMergeableGraph(std::shared_ptr<const GraphRepresentation> graph,
   std::string input_shape =
     iniparser_getstring(ini, (sec_name + ":Input_Shape").c_str(), "");
   if (!input_shape.empty()) {
-    g[0]->setProperty(Layer::PropertyType::input_shape, input_shape);
+    g[0]->getObject()->setProperty(Layer::PropertyType::input_shape,
+                                   input_shape);
   }
 
   std::string input_layers =
     iniparser_getstring(ini, (sec_name + ":Input_Layers").c_str(), "");
   if (!input_layers.empty()) {
-    g[0]->setProperty(Layer::PropertyType::input_layers, input_layers);
+    g[0]->getObject()->setProperty(Layer::PropertyType::input_layers,
+                                   input_layers);
   }
 
   return g;
@@ -271,7 +274,7 @@ void IniGraphInterpreter::serialize(
 
   std::vector<IniSection> sections;
   for (const auto &ln : representation->getSorted()) {
-    const auto &layer = ln.getObject();
+    const auto &layer = ln->getObject();
 
     IniSection s(layer->getName());
     s.setEntry("type", layer->getType());
@@ -333,7 +336,7 @@ IniGraphInterpreter::deserialize(const std::string &in) {
       }
       /** Parse all the layers defined as sections in order */
       ml_logd("probing section_name: %s", sec_name_);
-      std::shared_ptr<Layer> layer;
+      std::shared_ptr<LayerNode> layer;
 
       /**
        * If this section is a backbone, load backbone section from this
diff --git a/nntrainer/dataset/databuffer_func.h b/nntrainer/dataset/databuffer_func.h
index 7c5f93ca..c1712fd3 100644
--- a/nntrainer/dataset/databuffer_func.h
+++ b/nntrainer/dataset/databuffer_func.h
@@ -49,7 +49,7 @@ public:
   /**
    * @brief     Destructor
    */
-  ~DataBufferFromCallback(){};
+  ~DataBufferFromCallback() = default;
 
   /**
    * @brief     Initialize Buffer
diff --git a/nntrainer/graph/network_graph.cpp b/nntrainer/graph/network_graph.cpp
index c7949618..9dac4e1f 100644
--- a/nntrainer/graph/network_graph.cpp
+++ b/nntrainer/graph/network_graph.cpp
@@ -71,12 +71,13 @@ int NetworkGraph::compile(const LossType loss_type) {
    * Now that graph is compiled, remove all edges to save memory.
    * NodeList is kept for now for O(1) access of layers by idx.
    */
-  for (unsigned int i = 0; i < adj.size(); ++i)
+  for (unsigned int i = 0; i < adj.size(); ++i) {
     /**
      * As this resize is guaranteed to not insert new elements,  create a
      * default element needed by resize.
      */
-    adj[i].resize(1, LayerNode(nullptr, 0));
+    adj[i].resize(1);
+  }
 
   compiled = true;
 
@@ -86,7 +87,7 @@ int NetworkGraph::compile(const LossType loss_type) {
 void NetworkGraph::updateConnectionName(const std::string &from,
                                         const std::string &to) {
   for (unsigned int i = 0; i < adj.size(); ++i) {
-    auto &layer = adj[i].front().getObject();
+    auto &layer = adj[i].front()->getObject();
     if (istrequal(layer->getName(), to))
       continue;
     for (unsigned int j = 0; j < layer->input_layers.size(); ++j) {
@@ -99,8 +100,8 @@ void NetworkGraph::updateConnectionName(const std::string &from,
 
 void NetworkGraph::addDefaultInputLayers() {
   for (unsigned int i = 1; i < adj.size(); ++i) {
-    auto &layer = adj[i].front().getObject();
-    auto &prev_layer = adj[i - 1].front().getObject();
+    auto &layer = adj[i].front()->getObject();
+    auto &prev_layer = adj[i - 1].front()->getObject();
     if (layer->input_layers.size() == 0) {
       layer->input_layers.push_back(prev_layer->getName());
     }
@@ -108,39 +109,42 @@ void NetworkGraph::addDefaultInputLayers() {
 }
 
 void NetworkGraph::addLayerNode(std::shared_ptr<Layer> layer) {
-  std::list<LayerNode> l;
-  std::unique_ptr<LayerNode> node =
-    std::make_unique<LayerNode>(layer, adj.size());
+  addLayerNode(std::make_unique<LayerNode>(layer, adj.size()));
+}
 
-  l.assign(1, *node);
+void NetworkGraph::addLayerNode(std::shared_ptr<LayerNode> layer) {
+  std::list<std::shared_ptr<LayerNode>> l;
+  layer->setIndex(adj.size());
+
+  l.push_back(layer);
   adj.push_back(l);
 }
 
-LayerNode &NetworkGraph::getLayerNode(unsigned int ith) {
+std::shared_ptr<LayerNode> &NetworkGraph::getLayerNode(unsigned int ith) {
   if (ith >= size())
     throw std::invalid_argument("Exceed total number of layer");
 
-  if (adj[ith].front().getIndex() != ith)
+  if (adj[ith].front()->getIndex() != ith)
     throw std::runtime_error("Graph internal index mismatch");
 
   return adj[ith].front();
 }
 
-LayerNode &NetworkGraph::getSortedLayerNode(unsigned int ith) {
+std::shared_ptr<LayerNode> &NetworkGraph::getSortedLayerNode(unsigned int ith) {
   if (ith >= getSorted().size())
     throw std::invalid_argument("Exceed total number of layer");
 
   return getSorted()[ith];
 }
 
-void NetworkGraph::topologicalSortUtil(unsigned int ith,
-                                       std::vector<bool> &visited,
-                                       std::stack<LayerNode> &Stack) {
+void NetworkGraph::topologicalSortUtil(
+  unsigned int ith, std::vector<bool> &visited,
+  std::stack<std::shared_ptr<LayerNode>> &Stack) {
   visited[ith] = true;
 
-  std::list<LayerNode>::iterator i;
+  std::list<std::shared_ptr<LayerNode>>::iterator i;
   for (i = adj[ith].begin(); i != adj[ith].end(); ++i) {
-    auto index = (*i).getIndex();
+    auto index = (*i)->getIndex();
     if (!visited[index])
       topologicalSortUtil(index, visited, Stack);
   }
@@ -150,7 +154,7 @@ void NetworkGraph::topologicalSortUtil(unsigned int ith,
 
 void NetworkGraph::countNonTrainableLayersAtBegin() {
   for (auto iter = Sorted.cbegin(); iter != Sorted.cend(); iter++) {
-    if ((*iter).getObject()->getTrainable()) {
+    if ((*iter)->getObject()->getTrainable()) {
       skip_non_trainable_layers = iter - Sorted.cbegin();
       return;
     }
@@ -159,7 +163,7 @@ void NetworkGraph::countNonTrainableLayersAtBegin() {
 }
 
 void NetworkGraph::topologicalSort() {
-  std::stack<LayerNode> Stack;
+  std::stack<std::shared_ptr<LayerNode>> Stack;
   std::vector<bool> visited(adj.size());
   Sorted.clear();
 
@@ -367,11 +371,11 @@ int NetworkGraph::addLossLayer(const LossType loss_type) {
 
   int status = ML_ERROR_NONE;
 
-  if (Sorted.back().getObject()->getType() == LossLayer::type)
+  if (Sorted.back()->getObject()->getType() == LossLayer::type)
     return status;
 
-  if (Sorted.back().getObject()->getType() == TimeDistLayer::type) {
-    if (std::static_pointer_cast<TimeDistLayer>(Sorted.back().getObject())
+  if (Sorted.back()->getObject()->getType() == TimeDistLayer::type) {
+    if (std::static_pointer_cast<TimeDistLayer>(Sorted.back()->getObject())
           ->getDistLayerType() == LossLayer::type)
       return status;
   }
@@ -386,11 +390,11 @@ int NetworkGraph::addLossLayer(const LossType loss_type) {
     NN_RETURN_STATUS();
   }
 
-  LayerNode last_node = Sorted.back();
+  auto &last_node = Sorted.back();
   if (updated_loss_type == LossType::LOSS_ENTROPY) {
-    auto type = last_node.getObject()->getType();
+    auto type = last_node->getObject()->getType();
     if (type == TimeDistLayer::type) {
-      type = std::dynamic_pointer_cast<TimeDistLayer>(last_node.getObject())
+      type = std::dynamic_pointer_cast<TimeDistLayer>(last_node->getObject())
                ->getDistLayerType();
     }
 
@@ -400,11 +404,11 @@ int NetworkGraph::addLossLayer(const LossType loss_type) {
       return ML_ERROR_NOT_SUPPORTED;
     }
 
-    LayerNode last_node = Sorted.back();
+    /// TODO: the last entry in adj and sorted are not same
     adj.pop_back();
     Sorted.pop_back();
 
-    switch (last_node.getObject()->getActivationType()) {
+    switch (last_node->getObject()->getActivationType()) {
     case ActivationType::ACT_SIGMOID:
       updated_loss_type = LossType::LOSS_ENTROPY_SIGMOID;
       break;
@@ -417,27 +421,33 @@ int NetworkGraph::addLossLayer(const LossType loss_type) {
     }
   }
 
+  /// @note this assumes that loss layer only supports single input
   last_node = Sorted.back();
-  std::string input_str = last_node.getObject()->getName();
 
-  std::shared_ptr<Layer> layer = nntrainer::createLayer(LossLayer::type);
+  /// TODO: need to find all the nodes whose entry matches with node remove and
+  /// update them all
+  if (updated_loss_type == LossType::LOSS_ENTROPY_SIGMOID ||
+      updated_loss_type == LossType::LOSS_ENTROPY_SOFTMAX)
+    adj[last_node->getIndex()].resize(1);
 
+  std::shared_ptr<Layer> layer = nntrainer::createLayer(LossLayer::type);
   status =
     std::dynamic_pointer_cast<LossLayer>(layer)->setLoss(updated_loss_type);
   NN_RETURN_STATUS();
-
   ensureName(layer);
 
-  if (last_node.getObject()->getType() == TimeDistLayer::type) {
+  std::string input_str = last_node->getObject()->getName();
+
+  if (last_node->getObject()->getType() == TimeDistLayer::type) {
     std::string unit_str = layer->getName();
     ensureName(layer, "", "_unit");
     layer = distributeLayer(layer);
     layer->setName(unit_str);
   }
 
-  last_node.getObject()->setNumOutputs(1);
-  last_node.getObject()->output_layers.clear();
-  last_node.getObject()->output_layers.push_back(layer->getName());
+  last_node->getObject()->setNumOutputs(1);
+  last_node->getObject()->output_layers.clear();
+  last_node->getObject()->output_layers.push_back(layer->getName());
 
   layer->setNumInputs(1);
   layer->input_layers.clear();
@@ -465,9 +475,9 @@ void NetworkGraph::setOutputLayers() {
 
   size_t last_layer_count = 0;
   for (unsigned int idx = 0; idx < adj.size(); ++idx) {
-    auto &layer_idx = adj[idx].front().getObject();
+    auto &layer_idx = adj[idx].front()->getObject();
     for (unsigned int i = 0; i < adj.size(); ++i) {
-      auto &layer_i = adj[i].front().getObject();
+      auto &layer_i = adj[i].front()->getObject();
       if (istrequal(layer_i->getName(), layer_idx->getName()))
         continue;
       for (unsigned int j = 0; j < layer_i->input_layers.size(); ++j) {
@@ -511,7 +521,7 @@ void NetworkGraph::setOutputLayers() {
   }
 
   for (auto iter = adj.begin(); iter < adj.end(); ++iter) {
-    if ((*iter).front().getObject()->output_layers.size() == 0)
+    if ((*iter).front()->getObject()->output_layers.size() == 0)
       throw std::runtime_error("There is un-connected node");
   }
 }
@@ -531,7 +541,7 @@ int NetworkGraph::isCompilable() {
 }
 
 int NetworkGraph::checkCompiledGraph() {
-  auto &l = Sorted[0].getObject();
+  auto &l = Sorted[0]->getObject();
   /** First layer cannot be activation, batch normalization or loss */
   const std::string &type = l->getType();
   if (istrequal(type, ActivationLayer::type) ||
@@ -544,8 +554,8 @@ int NetworkGraph::checkCompiledGraph() {
 
   /** Dimension of input layers must be known */
   for (auto const &lnode : Sorted) {
-    if (lnode.getObject()->getType() == InputLayer::type) {
-      if (lnode.getObject()->getInputDimension().size() == 0) {
+    if (lnode->getObject()->getType() == InputLayer::type) {
+      if (lnode->getObject()->getInputDimension().size() == 0) {
         ml_loge("InputDimension of first layer is not set");
         return ML_ERROR_INVALID_PARAMETER;
       }
@@ -565,7 +575,7 @@ int NetworkGraph::realizeGraph() {
   /** This loop modifes adj. Get the size of adj preemptively. */
 
   for (unsigned int i = 0; i < adj_size_before_realize; ++i) {
-    Layer &l = *adj[i].front().getObject();
+    Layer &l = *adj[i].front()->getObject();
     ml_logd("layer name: %s", l.getName().c_str());
 
     /** If a layer does not has input nodes, then it must have input dimension
@@ -612,18 +622,20 @@ int NetworkGraph::realizeGraph() {
   return status;
 }
 
-LayerNode &NetworkGraph::getLayerNode(const std::string &layer_name) {
-
+std::shared_ptr<LayerNode> &
+NetworkGraph::getLayerNode(const std::string &layer_name) {
   for (auto &lnode_list : adj) {
     auto &lnode = lnode_list.front();
-    if (istrequal(lnode.getObject()->getName(), layer_name))
+    if (istrequal(lnode->getObject()->getName(), layer_name))
       return lnode;
   }
 
-  throw std::invalid_argument("Cannot find Layer");
+  std::stringstream ss;
+  ss << "Cannot find Layer: " << layer_name;
+  throw std::invalid_argument(ss.str());
 }
 
-void NetworkGraph::addEdge(unsigned int ith, LayerNode &node) {
+void NetworkGraph::addEdge(unsigned int ith, std::shared_ptr<LayerNode> &node) {
   if (ith >= adj.size())
     throw std::invalid_argument("Exceed total number of layer");
 
@@ -631,13 +643,13 @@ void NetworkGraph::addEdge(unsigned int ith, LayerNode &node) {
 }
 
 void NetworkGraph::connectGraph(unsigned int adj_idx) {
-  std::list<LayerNode>::iterator iter = adj[adj_idx].begin();
+  std::list<std::shared_ptr<LayerNode>>::iterator iter = adj[adj_idx].begin();
 
-  for (unsigned int j = 0; j < (*iter).getObject()->input_layers.size(); ++j) {
-    if (istrequal((*iter).getObject()->input_layers[j], "__data__"))
+  for (unsigned int j = 0; j < (*iter)->getObject()->input_layers.size(); ++j) {
+    if (istrequal((*iter)->getObject()->input_layers[j], "__data__"))
       continue;
     unsigned int to_node_id =
-      getLayerNode((*iter).getObject()->input_layers[j]).getIndex();
+      getLayerNode((*iter)->getObject()->input_layers[j])->getIndex();
     addEdge(to_node_id, (*iter));
   }
 }
@@ -652,19 +664,19 @@ int NetworkGraph::connectGraph() {
 
 void NetworkGraph::setBatchSize(unsigned int batch_size) {
   for (auto const &layer_adj_list : adj) {
-    layer_adj_list.front().getObject()->setBatch(batch_size);
+    layer_adj_list.front()->getObject()->setBatch(batch_size);
   }
 }
 
 sharedConstTensors NetworkGraph::forwarding(bool training) {
   for (auto const &ln : getSorted()) {
-    START_PROFILE(ln.event_key);
-    ln.getObject()->forwarding(training);
-    END_PROFILE(ln.event_key);
+    START_PROFILE(ln->event_key);
+    ln->getObject()->forwarding(training);
+    END_PROFILE(ln->event_key);
   }
 
   std::vector<sharedConstTensor> out;
-  for (auto const &nh : getSorted().back().getObject()->net_hidden)
+  for (auto const &nh : getSorted().back()->getObject()->net_hidden)
     out.push_back(MAKE_SHARED_TENSOR(nh->getVariable()));
 
   return out;
@@ -673,16 +685,16 @@ sharedConstTensors NetworkGraph::forwarding(bool training) {
 std::vector<TensorDim> NetworkGraph::getInputDimension() const {
   NNTR_THROW_IF(this->empty(), std::invalid_argument)
     << "[NetworkGraph] the graph has no node!";
-  return getSorted()[0].getObject()->getInputDimension();
+  return getSorted()[0]->getObject()->getInputDimension();
 }
 
 std::vector<TensorDim> NetworkGraph::getOutputDimension() const {
   NNTR_THROW_IF(this->empty(), std::invalid_argument)
     << "[NetworkGraph] the graph has no node!";
-  return getSorted().back().getObject()->getOutputDimension();
+  return getSorted().back()->getObject()->getOutputDimension();
 }
 
-std::vector<std::shared_ptr<Layer>>
+std::vector<std::shared_ptr<LayerNode>>
 NetworkGraph::getUnsortedLayers(const std::string &input_layer,
                                 const std::string &output_layer) const {
   /// @fixme: this won't work if input, output layers are not in order
@@ -693,7 +705,7 @@ NetworkGraph::getUnsortedLayers(const std::string &input_layer,
   unsigned int num_layers_remove_end = 0;
   if (!output_layer.empty()) {
     for (auto iter = adj.rbegin(); iter != adj.rend(); iter++) {
-      if ((*iter).front().getObject()->getName() != output_layer)
+      if ((*iter).front()->getObject()->getName() != output_layer)
         num_layers_remove_end++;
       else
         break;
@@ -708,7 +720,7 @@ NetworkGraph::getUnsortedLayers(const std::string &input_layer,
   if (!input_layer.empty()) {
     for (auto iter = adj.begin(); iter != adj.end() - num_layers_remove_end;
          iter++) {
-      if ((*iter).front().getObject()->getName() != input_layer)
+      if ((*iter).front()->getObject()->getName() != input_layer)
         num_layers_remove_start++;
       else
         break;
@@ -716,28 +728,28 @@ NetworkGraph::getUnsortedLayers(const std::string &input_layer,
   }
 
   /** copy the graph and return */
-  std::vector<std::shared_ptr<Layer>> ret;
+  std::vector<std::shared_ptr<LayerNode>> ret;
   std::transform(adj.begin() + num_layers_remove_start,
                  adj.end() - num_layers_remove_end, std::back_inserter(ret),
-                 [](auto const &elem) { return elem.front().getObject(); });
+                 [](auto const &elem) { return elem.front(); });
 
   return ret;
 }
 
-std::vector<std::shared_ptr<Layer>> NetworkGraph::getLayers() const {
-  std::vector<std::shared_ptr<Layer>> ret;
+std::vector<std::shared_ptr<LayerNode>> NetworkGraph::getLayerNodes() const {
+  std::vector<std::shared_ptr<LayerNode>> ret;
   if (compiled) {
     std::transform(Sorted.begin(), Sorted.end(), std::back_inserter(ret),
-                   [](auto const &elem) { return elem.getObject(); });
+                   [](auto const &elem) { return elem; });
   } else {
     std::transform(adj.begin(), adj.end(), std::back_inserter(ret),
-                   [](auto const &elem) { return elem.front().getObject(); });
+                   [](auto const &elem) { return elem.front(); });
   }
 
   return ret;
 }
 
-void NetworkGraph::extendGraph(std::vector<std::shared_ptr<Layer>> graph,
+void NetworkGraph::extendGraph(std::vector<std::shared_ptr<LayerNode>> graph,
                                std::string &prefix) {
 
   if (compiled)
@@ -750,21 +762,22 @@ void NetworkGraph::extendGraph(std::vector<std::shared_ptr<Layer>> graph,
    * This loop intends to connect a new backbone to be added with an old
    * backbone.
    */
-  for (unsigned int i = 0; i < graph[0]->input_layers.size(); ++i) {
-    if (sub_in_out.find(graph[0]->input_layers[i]) != sub_in_out.end()) {
-      graph[0]->input_layers[i] = sub_in_out[graph[0]->input_layers[i]];
-    } else if (layer_names.find(graph[0]->input_layers[i]) ==
-               layer_names.end()) {
+  auto &layer0_in = graph[0]->getObject()->input_layers;
+  for (unsigned int i = 0; i < layer0_in.size(); ++i) {
+    if (sub_in_out.find(layer0_in[i]) != sub_in_out.end()) {
+      layer0_in[i] = sub_in_out[layer0_in[i]];
+    } else if (layer_names.find(layer0_in[i]) == layer_names.end()) {
       throw std::runtime_error("Input layer name for backbone not found.");
     }
   }
 
   /** Insert the layer to the graph */
-  for (auto layer : graph) {
+  for (auto &layernode : graph) {
     /**
      * Add prefix to the existing layer name,
      * and ensure it is unique in this new graph
      */
+    auto &layer = layernode->getObject();
     std::string orig_name = prefix + layer->getName();
     ensureName(layer, prefix, "", true);
     sub_in_out.insert(std::make_pair(orig_name, layer->getName()));
@@ -779,20 +792,20 @@ void NetworkGraph::extendGraph(std::vector<std::shared_ptr<Layer>> graph,
       }
     }
 
-    addLayerNode(layer);
+    addLayerNode(layernode);
   }
 
   /** This allows connecting a layer to the backbone */
   sub_in_out.insert(
-    std::make_pair(prefix, adj.back().front().getObject()->getName()));
+    std::make_pair(prefix, adj.back().front()->getObject()->getName()));
 }
 
-void NetworkGraph::addLayer(std::shared_ptr<Layer> layer) {
+void NetworkGraph::addLayer(std::shared_ptr<LayerNode> layer) {
   if (compiled)
     throw std::runtime_error("Cannot modify graph after compile");
 
   /** Ensure that the layer has a name and is unique */
-  ensureName(layer);
+  ensureName(layer->getObject());
 
   /** Insert the layer to the graph */
   addLayerNode(layer);
@@ -801,7 +814,7 @@ void NetworkGraph::addLayer(std::shared_ptr<Layer> layer) {
 void NetworkGraph::inPlaceOptimize(const std::string &layer_type,
                                    Manager &manager) {
   for (auto &layer_node : getSorted()) {
-    auto &l = layer_node.getObject();
+    auto &l = layer_node->getObject();
     std::string l_type = l->getType();
     if (l_type == TimeDistLayer::type) {
       l_type = std::dynamic_pointer_cast<TimeDistLayer>(l)->getDistLayerType();
@@ -814,7 +827,7 @@ void NetworkGraph::inPlaceOptimize(const std::string &layer_type,
       if (l->input_layers.size() != 1)
         throw std::runtime_error("Internal error in the formed graph");
 
-      auto &prev_layer = getLayerNode(l->input_layers[0]).getObject();
+      auto &prev_layer = getLayerNode(l->input_layers[0])->getObject();
 
       unsigned int loc;
       auto layer_name = l->getName();
@@ -897,14 +910,14 @@ void NetworkGraph::inPlaceOptimize(Manager &manager) {
     inPlaceOptimize(layer_type, manager);
 }
 
-const std::vector<LayerNode> &NetworkGraph::getSorted() const {
+const std::vector<std::shared_ptr<LayerNode>> &NetworkGraph::getSorted() const {
   if (!compiled)
     throw std::runtime_error("Cannot get sorted graph before compiling graph");
 
   return Sorted;
 }
 
-std::vector<LayerNode> &NetworkGraph::getSorted() {
+std::vector<std::shared_ptr<LayerNode>> &NetworkGraph::getSorted() {
   if (!compiled)
     throw std::runtime_error("Cannot get sorted graph before compiling graph");
 
diff --git a/nntrainer/graph/network_graph.h b/nntrainer/graph/network_graph.h
index fa020eb0..febc678f 100644
--- a/nntrainer/graph/network_graph.h
+++ b/nntrainer/graph/network_graph.h
@@ -54,7 +54,7 @@ public:
    * @brief Create new LayerNode and add into Graph
    * @param[in] layer shared_ptr of Layer
    */
-  void addLayer(std::shared_ptr<Layer> layer);
+  void addLayer(std::shared_ptr<LayerNode> layer);
 
   /**
    * @brief get current flat graph from the model before sorting
@@ -65,7 +65,7 @@ public:
    * @todo remove getting unsorted layers from model loader, compile model
    * loader
    */
-  std::vector<std::shared_ptr<Layer>>
+  std::vector<std::shared_ptr<LayerNode>>
   getUnsortedLayers(const std::string &input_layer,
                     const std::string &output_layer) const;
 
@@ -108,6 +108,7 @@ public:
    * @brief     reset the graph
    */
   void reset() {
+
     adj.clear();
     Sorted.clear();
     layer_names.clear();
@@ -120,21 +121,21 @@ public:
    * @param[in] index
    * @ret LayerNode
    */
-  LayerNode &getLayerNode(unsigned int ith);
+  std::shared_ptr<LayerNode> &getLayerNode(unsigned int ith);
 
   /**
    * @brief getter of Sorted LayerNode with index number
    * @param[in] index
    * @ret LayerNode
    */
-  LayerNode &getSortedLayerNode(unsigned int ith);
+  std::shared_ptr<LayerNode> &getSortedLayerNode(unsigned int ith);
 
   /**
    * @brief getter of LayerNode with layer name
    * @param[in] layer name
    * @retval LayerNode
    */
-  LayerNode &getLayerNode(const std::string &layer_name);
+  std::shared_ptr<LayerNode> &getLayerNode(const std::string &layer_name);
 
   /**
    * @brief getter of Layer with layer name
@@ -142,16 +143,16 @@ public:
    * @retval Layer
    */
   std::shared_ptr<Layer> getLayer(const std::string &layer_name) {
-    return getLayerNode(layer_name).getObject();
+    return getLayerNode(layer_name)->getObject();
   }
 
   /**
-   * @brief getter all the layers in the model
-   * @retval Layers
-   * @note these layers will be in sorted order if the model is compiled,
-   * otherwise the order is the order of addition of layers in the model.
+   * @brief getter all the layer nodes in the model
+   * @retval Layer nodes
+   * @note these layer nodes will be in sorted order if the model is compiled,
+   * otherwise the order is the order of addition of layer nodes in the model.
    */
-  std::vector<std::shared_ptr<Layer>> getLayers() const;
+  std::vector<std::shared_ptr<LayerNode>> getLayerNodes() const;
 
   /**
    * @brief     join passed graph into the existing graph model
@@ -163,7 +164,7 @@ public:
    *
    * @todo rename to addLayers
    */
-  void extendGraph(std::vector<std::shared_ptr<Layer>> graph,
+  void extendGraph(std::vector<std::shared_ptr<LayerNode>> graph,
                    std::string &prefix);
 
   /**
@@ -184,19 +185,20 @@ public:
    * @brief     getter of ordered graph
    * @retval    ordered LayerNode list
    */
-  const std::vector<LayerNode> &getSorted() const;
+  const std::vector<std::shared_ptr<LayerNode>> &getSorted() const;
 
   /**
    * @brief     getter of ordered graph
    * @retval    ordered LayerNode list
    */
-  std::vector<LayerNode> &getSorted();
+  std::vector<std::shared_ptr<LayerNode>> &getSorted();
 
   /**
    * @brief     get begin iterator for the backwarding
    * @retval    const reverse iterator marking the begin of backwarding
    */
-  std::vector<LayerNode>::const_reverse_iterator getBackwardingBeginIter() {
+  std::vector<std::shared_ptr<LayerNode>>::const_reverse_iterator
+  getBackwardingBeginIter() {
     return Sorted.crbegin();
   }
 
@@ -204,7 +206,8 @@ public:
    * @brief     get end iterator for the backwarding
    * @retval    const reverse iterator marking the end of backwarding
    */
-  std::vector<LayerNode>::const_reverse_iterator getBackwardingEndIter() {
+  std::vector<std::shared_ptr<LayerNode>>::const_reverse_iterator
+  getBackwardingEndIter() {
     return Sorted.crend() - skip_non_trainable_layers;
   }
 
@@ -234,7 +237,7 @@ public:
     if (this != &from) {
       // FIXME: this assumes elements already in layers/adj, solve that
       for (unsigned int i = 0; i < adj.size(); i++)
-        adj[i].front().getObject()->copy(from.adj[i].front().getObject());
+        adj[i].front()->getObject()->copy(from.adj[i].front()->getObject());
     }
     return *this;
   }
@@ -242,8 +245,10 @@ public:
 private:
   std::map<std::string, std::string> sub_in_out; /** This is map to identify
                    input and output layer name of subgraph */
-  std::vector<std::list<LayerNode>> adj;         /**< Graph Structure */
-  std::vector<LayerNode> Sorted; /**< Ordered Graph Node List  */
+  std::vector<std::list<std::shared_ptr<LayerNode>>>
+    adj; /**< Graph Structure */
+  std::vector<std::shared_ptr<LayerNode>>
+    Sorted; /**< Ordered Graph Node List  */
   std::set<std::string>
     layer_names; /**< Set containing all the names of layers in the model */
   int def_name_count; /**< Count assigned to layer names declared by default */
@@ -259,7 +264,7 @@ private:
    * @param[in] stack for Node list to visit.
    */
   void topologicalSortUtil(unsigned int ith, std::vector<bool> &visited,
-                           std::stack<LayerNode> &Stack);
+                           std::stack<std::shared_ptr<LayerNode>> &Stack);
 
   /**
    * @brief     check if graph is ready to compile.
@@ -280,7 +285,7 @@ private:
    * @param[in] ith Node index : From
    * @param[in] node LayerNode object to be added : To
    */
-  void addEdge(unsigned int ith, LayerNode &node);
+  void addEdge(unsigned int ith, std::shared_ptr<LayerNode> &node);
 
   /**
    * @brief     make connection between nodes
@@ -371,6 +376,12 @@ private:
    */
   void addLayerNode(std::shared_ptr<Layer> layer);
 
+  /**
+   * @brief Add given LayerNode to the Graph
+   * @param[in] layer shared_ptr of LayerNode
+   */
+  void addLayerNode(std::shared_ptr<LayerNode> layer);
+
   /**
    * @brief Sorting and Define order to calculate : Depth First Search
    */
diff --git a/nntrainer/layers/layer_internal.h b/nntrainer/layers/layer_internal.h
index 4ae88ed5..ce7a24e1 100644
--- a/nntrainer/layers/layer_internal.h
+++ b/nntrainer/layers/layer_internal.h
@@ -43,8 +43,11 @@ namespace nntrainer {
 /**
  * @class   Layer Base class for layers
  * @brief   Base class for all layers
+ *
+ * @details nntrainer::Layer inherits ml::train::Layer but has been ommitted to
+ * disallow static_cast between nntrainer::Layer and ml::train::Layer objects.
  */
-class Layer : public ml::train::Layer {
+class Layer {
 
   /** model classes can call private methods which arent exposed to public */
   friend class NeuralNetwork;
@@ -76,6 +79,11 @@ public:
     setNumOutputs(1);
   }
 
+  /**
+   * @brief     Destructor of Layer Class
+   */
+  virtual ~Layer() = default;
+
   /**
    *  @brief  Move constructor of Layer.
    *  @param[in] Layer &&
@@ -88,6 +96,12 @@ public:
    */
   virtual Layer &operator=(Layer &&rhs) = default;
 
+  /**
+   * @brief Get the layer type
+   * @return const std::string type representation
+   */
+  virtual const std::string getType() const = 0;
+
   /**
    * @brief     Forward Propagation of a layer
    * @param[in] in List of Input Tensors taken by this layer
@@ -711,8 +725,8 @@ std::ostream &operator<<(std::ostream &out, T &l) {
   return out;
 }
 
-using CreateLayerFunc = ml::train::Layer *(*)();
-using DestroyLayerFunc = void (*)(ml::train::Layer *);
+using CreateLayerFunc = nntrainer::Layer *(*)();
+using DestroyLayerFunc = void (*)(nntrainer::Layer *);
 
 /**
  * @brief  Layer Pluggable struct that enables pluggable layer
@@ -728,6 +742,23 @@ typedef struct {
  */
 extern "C" LayerPluggable ml_train_layer_pluggable;
 
+/**
+ * @brief General Layer Factory function to register Layer
+ *
+ * @param props property representation
+ * @return std::unique_ptr<ml::train::Layer> created object
+ */
+template <typename T,
+          std::enable_if_t<std::is_base_of<Layer, T>::value, T> * = nullptr>
+std::unique_ptr<Layer> createLayer(const std::vector<std::string> &props = {}) {
+  std::unique_ptr<Layer> ptr = std::make_unique<T>();
+
+  if (ptr->setProperty(props) != ML_ERROR_NONE) {
+    throw std::invalid_argument("Set properties failed for layer");
+  }
+  return ptr;
+}
+
 } // namespace nntrainer
 
 #endif /* __cplusplus */
diff --git a/nntrainer/layers/layer_node.cpp b/nntrainer/layers/layer_node.cpp
new file mode 100644
index 00000000..1b85bc94
--- /dev/null
+++ b/nntrainer/layers/layer_node.cpp
@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: Apache-2.0
+/**
+ * Copyright (C) 2021 Parichay Kapoor <pk.kapoor@samsung.com>
+ *
+ * @file   layer_node.cpp
+ * @date   1 April 2021
+ * @see    https://github.com/nnstreamer/nntrainer
+ * @author Parichay Kapoor <pk.kapoor@samsung.com>
+ * @bug    No known bugs except for NYI items
+ * @brief  This is the layer node for network graph
+ */
+
+#include <layer_node.h>
+
+namespace nntrainer {
+
+std::shared_ptr<Layer> getLayerDevel(std::shared_ptr<ml::train::Layer> l) {
+  std::shared_ptr<LayerNode> lnode = std::static_pointer_cast<LayerNode>(l);
+
+  std::shared_ptr<Layer> &layer = lnode->getObject();
+
+  return layer;
+}
+
+}; // namespace nntrainer
diff --git a/nntrainer/layers/layer_node.h b/nntrainer/layers/layer_node.h
index 4fee32b7..f62fe61b 100644
--- a/nntrainer/layers/layer_node.h
+++ b/nntrainer/layers/layer_node.h
@@ -2,12 +2,12 @@
 /**
  * Copyright (C) 2021 Parichay Kapoor <pk.kapoor@samsung.com>
  *
- * @file   graph_node.h
+ * @file   layer_node.h
  * @date   1 April 2021
  * @see    https://github.com/nnstreamer/nntrainer
  * @author Parichay Kapoor <pk.kapoor@samsung.com>
  * @bug    No known bugs except for NYI items
- * @brief  This is the graph node interface for c++ API
+ * @brief  This is the layer node for network graph
  */
 
 #ifndef __LAYER_NODE_H__
@@ -25,11 +25,16 @@ namespace nntrainer {
  */
 class LayerNode : public ml::train::Layer, public GraphNode {
 public:
+  /**
+   * @brief     Default constructor
+   */
+  LayerNode() : LayerNode(nullptr) {}
+
   /**
    * @brief     Constructor of LayerNode class
    *
    */
-  LayerNode(std::shared_ptr<nntrainer::Layer> l, size_t idx) :
+  LayerNode(std::shared_ptr<nntrainer::Layer> l, size_t idx = 0) :
     layer(l),
     index(idx),
     flatten(false),
@@ -41,6 +46,11 @@ public:
    */
   ~LayerNode() = default;
 
+  /**
+   * @brief     Set the index for the node
+   */
+  void setIndex(size_t idx) { index = idx; }
+
   /**
    * Support all the interface requirements by ml::train::Layer
    */
@@ -109,6 +119,7 @@ public:
 #endif
 
 private:
+  // TODO: make this unique_ptr once getObject API is removed
   std::shared_ptr<nntrainer::Layer>
     layer;      /**< The actual object in the graph node */
   size_t index; /**< index of each node */
@@ -120,5 +131,13 @@ private:
     activation_type; /**< activation applied to the output of this node */
 };
 
+/**
+ * @brief   Get Layer devel from ml::train::Layer
+ *
+ * @param   l Layer object
+ * @return  Layer devel object
+ */
+std::shared_ptr<Layer> getLayerDevel(std::shared_ptr<ml::train::Layer> l);
+
 } // namespace nntrainer
 #endif // __LAYER_NODE_H__
diff --git a/nntrainer/layers/meson.build b/nntrainer/layers/meson.build
index 7ebd68fc..21416f49 100644
--- a/nntrainer/layers/meson.build
+++ b/nntrainer/layers/meson.build
@@ -11,6 +11,7 @@ layer_sources = [
   'output_layer.cpp',
   'layer.cpp',
   'layer_factory.cpp',
+  'layer_node.cpp',
   'loss_layer.cpp',
   'pooling2d_layer.cpp',
   'preprocess_flip_layer.cpp',
diff --git a/nntrainer/layers/plugged_layer.h b/nntrainer/layers/plugged_layer.h
index 02865f2a..75ebfb30 100644
--- a/nntrainer/layers/plugged_layer.h
+++ b/nntrainer/layers/plugged_layer.h
@@ -37,7 +37,7 @@ public:
   PluggedLayer(const nntrainer::LayerPluggable *pluggable) :
     /// @todo we won't need dynamic pointer cast here after api is fully
     /// implemented
-    layerImpl(dynamic_cast<nntrainer::Layer *>(pluggable->createfunc())),
+    layerImpl(pluggable->createfunc()),
     destroy_func(pluggable->destroyfunc) {
     NNTR_THROW_IF(layerImpl == nullptr, std::invalid_argument)
       << "either create_func_ failed or cannot dynamic cast to layer_internal";
diff --git a/nntrainer/models/neuralnet.cpp b/nntrainer/models/neuralnet.cpp
index 458c41af..766ef91b 100644
--- a/nntrainer/models/neuralnet.cpp
+++ b/nntrainer/models/neuralnet.cpp
@@ -171,13 +171,14 @@ int NeuralNetwork::initialize() {
   for (unsigned int idx = 0; idx < n_layers; ++idx) {
     bool first = idx == 0;
     auto &lnode = model_graph.getSortedLayerNode(idx);
-    Layer &l = *lnode.getObject();
-    ml_logd("layer name : %s", l.getName().c_str());
+    auto &lptr = lnode->getObject();
+    ml_logd("layer name : %s", lptr->getName().c_str());
     std::string cur_type;
-    if (l.getType() == TimeDistLayer::type) {
-      cur_type = dynamic_cast<TimeDistLayer &>(l).getDistLayerType();
+    if (lptr->getType() == TimeDistLayer::type) {
+      cur_type =
+        std::dynamic_pointer_cast<TimeDistLayer>(lptr)->getDistLayerType();
     } else {
-      cur_type = l.getType();
+      cur_type = lptr->getType();
     }
 
     /**
@@ -186,10 +187,10 @@ int NeuralNetwork::initialize() {
      */
     if (!first) {
       std::string l_pre_type =
-        model_graph.getSortedLayerNode(idx - 1).getObject()->getType();
+        model_graph.getSortedLayerNode(idx - 1)->getObject()->getType();
       if (l_pre_type == TimeDistLayer::type) {
         l_pre_type = std::dynamic_pointer_cast<TimeDistLayer>(
-                       model_graph.getSortedLayerNode(idx - 1).getObject())
+                       model_graph.getSortedLayerNode(idx - 1)->getObject())
                        ->getDistLayerType();
       }
       if (istrequal(l_pre_type, ActivationLayer::type) &&
@@ -198,19 +199,19 @@ int NeuralNetwork::initialize() {
         return ML_ERROR_INVALID_PARAMETER;
       }
 
-      for (unsigned int i = 0; i < l.input_layers.size(); ++i) {
+      for (unsigned int i = 0; i < lptr->input_layers.size(); ++i) {
         Layer &in_layer =
-          *model_graph.getLayerNode(l.input_layers[i]).getObject();
+          *model_graph.getLayerNode(lptr->input_layers[i])->getObject();
 
         unsigned int location = 0;
         for (unsigned int j = 0; j < in_layer.output_layers.size(); ++j) {
-          if (in_layer.output_layers[j] == l.getName()) {
+          if (in_layer.output_layers[j] == lptr->getName()) {
             location = j;
             break;
           }
         }
 
-        l.setInputDimension(in_layer.getOutputDimension()[location], i);
+        lptr->setInputDimension(in_layer.getOutputDimension()[location], i);
       }
     }
 
@@ -218,38 +219,40 @@ int NeuralNetwork::initialize() {
      * Initialize all the layers, allocate output tensors for each layer
      * and add optimizer related weights for the layer
      */
-    status = l.initialize(*manager);
+    status = lptr->initialize(*manager);
     NN_RETURN_STATUS();
 
-    REGISTER_EVENT(l.getName(), lnode.event_key)
+    REGISTER_EVENT(lptr->getName(), lnode->event_key)
 
-    auto &in_out = manager->trackLayerOutputs(
-      cur_type, l.getName(), l.getOutputDimension(), l.getInputDimension());
-    l.setOutputBuffers(in_out);
+    auto &in_out = manager->trackLayerOutputs(cur_type, lptr->getName(),
+                                              lptr->getOutputDimension(),
+                                              lptr->getInputDimension());
+    lptr->setOutputBuffers(in_out);
 
     /** Connect the output of the previous layers with the input of the current
      * layer */
     if (!first) {
-      for (unsigned int i = 0; i < l.input_layers.size(); ++i) {
+      for (unsigned int i = 0; i < lptr->input_layers.size(); ++i) {
         Layer &in_layer =
-          *model_graph.getLayerNode(l.input_layers[i]).getObject();
+          *model_graph.getLayerNode(lptr->input_layers[i])->getObject();
 
         unsigned int location = 0;
         for (unsigned int j = 0; j < in_layer.output_layers.size(); ++j) {
-          if (in_layer.output_layers[j] == l.getName()) {
+          if (in_layer.output_layers[j] == lptr->getName()) {
             location = j;
             break;
           }
         }
 
-        l.net_input[i] = model_graph.getLayerNode(l.input_layers[i])
-                           .getObject()
-                           ->net_hidden[location];
+        lptr->net_input[i] = model_graph.getLayerNode(lptr->input_layers[i])
+                               ->getObject()
+                               ->net_hidden[location];
       }
     } else {
-      auto &in_out = manager->trackLayerInputs(
-        cur_type, l.getName(), l.getInputDimension(), l.getOutputDimension());
-      l.setInputBuffers(in_out);
+      auto &in_out = manager->trackLayerInputs(cur_type, lptr->getName(),
+                                               lptr->getInputDimension(),
+                                               lptr->getOutputDimension());
+      lptr->setInputBuffers(in_out);
     }
   }
 
@@ -258,7 +261,7 @@ int NeuralNetwork::initialize() {
     opt->initialize();
     for (unsigned int idx = 0; idx < n_layers; ++idx) {
       auto &lnode = model_graph.getSortedLayerNode(idx);
-      opt->addOptimizerVariable(lnode.getObject()->getWeightsRef());
+      opt->addOptimizerVariable(lnode->getObject()->getWeightsRef());
     }
   }
 
@@ -303,16 +306,16 @@ sharedConstTensors NeuralNetwork::forwarding(sharedConstTensors input,
       (!label.empty() && label[0]->batch() != batch_size))
     throw std::logic_error("Error: mismatch in batchsize for data and model.");
 
-  auto &first_layer = model_graph.getSortedLayerNode(0).getObject();
+  auto &first_layer = model_graph.getSortedLayerNode(0)->getObject();
   auto &last_layer =
     model_graph.getSortedLayerNode(model_graph.getSorted().size() - 1)
-      .getObject();
+      ->getObject();
 
   /// @note centroid_knn layer needs to be the last layer, currently it is
   /// not possible because loss layer is always added.
   /// if centroid_knn layer can be last layer, this loop is not required
   for (auto &layer_node : model_graph.getSorted()) {
-    auto l = layer_node.getObject();
+    auto &l = layer_node->getObject();
     if (l->getType() == "centroid_knn") {
       l->net_hidden[0]->getGradientRef() = *label[0].get();
     }
@@ -380,10 +383,11 @@ void NeuralNetwork::backwarding(int iteration) {
   auto iter_begin = model_graph.getBackwardingBeginIter();
   auto iter_end = model_graph.getBackwardingEndIter();
 
-  if (iter_begin->getObject()->getType() != LossLayer::type) {
+  auto &lptr_begin = (*iter_begin);
+  if (lptr_begin->getObject()->getType() != LossLayer::type) {
     bool has_loss = false;
-    if (iter_begin->getObject()->getType() == TimeDistLayer::type) {
-      if (std::dynamic_pointer_cast<TimeDistLayer>(iter_begin->getObject())
+    if (lptr_begin->getObject()->getType() == TimeDistLayer::type) {
+      if (std::dynamic_pointer_cast<TimeDistLayer>(lptr_begin->getObject())
             ->getDistLayerType() == LossLayer::type)
         has_loss = true;
     }
@@ -392,10 +396,10 @@ void NeuralNetwork::backwarding(int iteration) {
   }
 
   for (auto iter = iter_begin; iter != iter_end - 1; iter++) {
-    backwarding(iter->getObject(), iteration, true);
+    backwarding((*iter)->getObject(), iteration, true);
   }
 
-  auto last_layer = (iter_end - 1)->getObject();
+  auto last_layer = (*(iter_end - 1))->getObject();
   /**
    * The last trainable layer need not calculate the derivatives
    */
@@ -414,7 +418,7 @@ void NeuralNetwork::backwarding(int iteration) {
 void NeuralNetwork::backwarding(sharedConstTensors label, int iteration) {
   auto &loss_layer =
     model_graph.getSortedLayerNode(model_graph.getSorted().size() - 1)
-      .getObject();
+      ->getObject();
   loss_layer->net_hidden[0]->getGradientRef() = *label[0].get();
 
   backwarding(iteration);
@@ -425,7 +429,7 @@ float NeuralNetwork::getLoss() {
 
   auto &sorted = model_graph.getSorted();
   for (unsigned int i = 0; i < sorted.size(); i++) {
-    loss += sorted[i].getObject()->getLoss();
+    loss += sorted[i]->getObject()->getLoss();
   }
   return loss;
 }
@@ -467,7 +471,7 @@ void NeuralNetwork::saveModel() {
 
   auto &layers = model_graph.getSorted();
   for (unsigned int i = 0; i < layers.size(); i++)
-    layers[i].getObject()->save(model_file);
+    layers[i]->getObject()->save(model_file);
   model_file.write((char *)&epoch_idx, sizeof(epoch_idx));
   model_file.write((char *)&iter, sizeof(iter));
   model_file.close();
@@ -500,7 +504,7 @@ void NeuralNetwork::readModel() {
 
   auto &layers = tmp.model_graph.getSorted();
   for (unsigned int i = 0; i < layers.size(); i++)
-    layers[i].getObject()->read(model_file);
+    layers[i]->getObject()->read(model_file);
   checkedRead(model_file, (char *)&tmp.epoch_idx, sizeof(epoch_idx),
               "[NeuralNetwork::readModel] failed to read epoch_idx");
   checkedRead(model_file, (char *)&tmp.iter, sizeof(iter),
@@ -524,7 +528,7 @@ void NeuralNetwork::setBatchSize(unsigned int batch) {
 
 bool NeuralNetwork::validateInput(sharedConstTensors X) {
 
-  auto &first_layer = model_graph.getSortedLayerNode(0).getObject();
+  auto &first_layer = model_graph.getSortedLayerNode(0)->getObject();
   auto input_dim = first_layer->getInputDimension();
   if (X.size() != input_dim.size()) {
     ml_loge("Error: provided number of inputs %d, required %d", (int)X.size(),
@@ -574,7 +578,7 @@ sharedConstTensors NeuralNetwork::inference(sharedConstTensors X,
     return out;
   }
 
-  auto &last_layer = model_graph.getSorted().back().getObject();
+  auto &last_layer = model_graph.getSorted().back()->getObject();
   for (unsigned int i = 0; i < last_layer->getNumOutputs(); ++i) {
     out.push_back(MAKE_SHARED_TENSOR(last_layer->net_hidden[i]->getVariable()));
   }
@@ -658,10 +662,10 @@ int NeuralNetwork::train_run() {
     iter = 0;
   }
 
-  auto &first_layer = model_graph.getSortedLayerNode(0).getObject();
+  auto &first_layer = model_graph.getSortedLayerNode(0)->getObject();
   auto &last_layer =
     model_graph.getSortedLayerNode(model_graph.getSorted().size() - 1)
-      .getObject();
+      ->getObject();
 
   auto &output = last_layer->net_hidden[0]->getVariableRef();
   auto &label = last_layer->net_hidden[0]->getGradientRef();
@@ -670,7 +674,7 @@ int NeuralNetwork::train_run() {
   /// @todo migrate this to trait based system; sth like need label?
   std::shared_ptr<Layer> layer_;
   for (auto &layer_node : model_graph.getSorted()) {
-    layer_ = layer_node.getObject();
+    layer_ = layer_node->getObject();
     if (layer_->getType() == "centroid_knn") {
       layer_->net_hidden[0]->getGradientRef() = label;
     }
@@ -778,7 +782,7 @@ int NeuralNetwork::addLayer(NodeType layer) {
   }
 
   /** Validate the layer to be added */
-  status = layer->checkValidation();
+  status = layer->getObject()->checkValidation();
   if (status != ML_ERROR_NONE) {
     ml_loge("layer(%s) validation failed.", layer->getName().c_str());
     return status;
@@ -802,7 +806,7 @@ int NeuralNetwork::extendGraph(GraphType graph, std::string prefix) {
   return ML_ERROR_NONE;
 }
 
-std::vector<std::shared_ptr<Layer>>
+NeuralNetwork::GraphType
 NeuralNetwork::getUnsortedLayers(const std::string &input_layer,
                                  const std::string &output_layer) {
   return model_graph.getUnsortedLayers(input_layer, output_layer);
@@ -827,20 +831,8 @@ int NeuralNetwork::setDataBuffer(std::shared_ptr<DataBuffer> data_buffer) {
 
 int NeuralNetwork::getLayer(const char *name,
                             std::shared_ptr<ml::train::Layer> *layer) {
-  std::shared_ptr<Layer> layer_;
-  int ret = getLayer(name, &layer_);
-  if (ret == ML_ERROR_NONE)
-    *layer = layer_;
-  return ret;
-}
-
-int NeuralNetwork::getLayer(const char *name, NodeType *layer) {
-  NodeType ret = model_graph.getLayer(std::string(name));
-
-  if (ret == nullptr)
-    return ML_ERROR_INVALID_PARAMETER;
-
-  *layer = ret;
+  *layer = std::static_pointer_cast<ml::train::Layer>(
+    model_graph.getLayerNode(std::string(name)));
   return ML_ERROR_NONE;
 }
 
@@ -911,7 +903,7 @@ void NeuralNetwork::print(std::ostream &out, unsigned int flags,
   }
 
   // TODO: get sorted layers if initialized
-  auto layers = model_graph.getLayers();
+  auto layers = model_graph.getLayerNodes();
   if (flags & PRINT_GRAPH_INFO) {
     out << "graph contains " << layers.size() << " operation nodes\n";
     /// @todo print graph info
@@ -940,7 +932,7 @@ void NeuralNetwork::print(std::ostream &out, unsigned int flags,
 
   /** print layer properties */
   for (auto &layer : layers)
-    layer->printPreset(out, layerPrintPreset);
+    layer->getObject()->printPreset(out, layerPrintPreset);
 
   /// @todo Add status to check neuralnet has been run. #290
 }
diff --git a/nntrainer/models/neuralnet.h b/nntrainer/models/neuralnet.h
index b4b18786..497bf11c 100644
--- a/nntrainer/models/neuralnet.h
+++ b/nntrainer/models/neuralnet.h
@@ -40,7 +40,7 @@
 #include <fc_layer.h>
 #include <flatten_layer.h>
 #include <input_layer.h>
-#include <layer_internal.h>
+#include <layer_node.h>
 #include <loss_layer.h>
 #include <manager.h>
 #include <ml-api-common.h>
@@ -78,8 +78,8 @@ class NeuralNetwork : public ml::train::Model {
   friend class ModelLoader; /** access private members of ModelLoader */
 
 public:
-  using NodeType = std::shared_ptr<Layer>; /** Type of a Node */
-  using GraphType = std::vector<NodeType>; /** actual graph type */
+  using NodeType = std::shared_ptr<LayerNode>; /** Type of a Node */
+  using GraphType = std::vector<NodeType>;     /** actual graph type */
   using FlatGraphType =
     std::vector<NodeType>; /** topological sorted, iterable 1-D list of nodes */
   using NetworkGraphType = nntrainer::NetworkGraph;
@@ -301,7 +301,7 @@ public:
    * @retval #ML_ERROR_INVALID_PARAMETER invalid parameter.
    */
   int addLayer(std::shared_ptr<ml::train::Layer> layer) {
-    return addLayer(std::static_pointer_cast<Layer>(layer));
+    return addLayer(std::static_pointer_cast<LayerNode>(layer));
   }
 
   /**
@@ -371,7 +371,7 @@ public:
    * @note these layers will be in sorted order if the model is compiled,
    * otherwise the order is the order of addition of layers in the model.
    */
-  FlatGraphType getFlatGraph() { return model_graph.getLayers(); }
+  FlatGraphType getFlatGraph() { return model_graph.getLayerNodes(); }
 
   /**
    * @brief     get network graph
diff --git a/test/unittest/compiler/unittest_interpreter.cpp b/test/unittest/compiler/unittest_interpreter.cpp
index 9ad00eda..aae3323e 100644
--- a/test/unittest/compiler/unittest_interpreter.cpp
+++ b/test/unittest/compiler/unittest_interpreter.cpp
@@ -34,9 +34,12 @@ makeGraph(const std::vector<LayerReprentation> &layer_reps) {
   auto graph = std::make_shared<nntrainer::GraphRepresentation>();
 
   for (const auto &layer_representation : layer_reps) {
-    std::shared_ptr<ml::train::Layer> layer = ac.createObject<ml::train::Layer>(
-      layer_representation.first, layer_representation.second);
-    graph->addLayer(std::static_pointer_cast<nntrainer::Layer>(layer));
+    std::shared_ptr<nntrainer::Layer> nntr_layer =
+      ac.createObject<nntrainer::Layer>(layer_representation.first,
+                                        layer_representation.second);
+    std::shared_ptr<nntrainer::LayerNode> layer =
+      std::make_unique<nntrainer::LayerNode>(nntr_layer);
+    graph->addLayer(layer);
   }
 
   return graph;
@@ -59,8 +62,8 @@ auto ini_interpreter =
  */
 static void graphEqual(const nntrainer::GraphRepresentation &lhs,
                        const nntrainer::GraphRepresentation &rhs) {
-  auto layers = lhs.getLayers();
-  auto ref_layers = rhs.getLayers();
+  const auto &layers = lhs.getLayerNodes();
+  const auto &ref_layers = rhs.getLayerNodes();
   EXPECT_EQ(layers.size(), ref_layers.size());
 
   auto is_node_equal = [](const nntrainer::Layer &l,
@@ -79,7 +82,7 @@ static void graphEqual(const nntrainer::GraphRepresentation &lhs,
 
   if (layers.size() == ref_layers.size()) {
     for (unsigned int i = 0; i < layers.size(); ++i) {
-      is_node_equal(*layers[i], *ref_layers[i]);
+      is_node_equal(*layers[i]->getObject(), *ref_layers[i]->getObject());
     }
   }
 }
diff --git a/test/unittest/unittest_nntrainer_modelfile.cpp b/test/unittest/unittest_nntrainer_modelfile.cpp
index 8857a7c9..12d77ec5 100644
--- a/test/unittest/unittest_nntrainer_modelfile.cpp
+++ b/test/unittest/unittest_nntrainer_modelfile.cpp
@@ -454,14 +454,16 @@ TEST(nntrainerIniTest, backbone_p_05) {
   EXPECT_EQ(flat_backbone.size(), flat_direct.size());
 
   for (size_t idx = 0; idx < flat_backbone.size(); idx++) {
-    EXPECT_EQ(flat_backbone[idx]->getType(), flat_direct[idx]->getType());
-    EXPECT_EQ(flat_backbone[idx]->getInputDimension(),
-              flat_direct[idx]->getInputDimension());
-    EXPECT_EQ(flat_backbone[idx]->getOutputDimension(),
-              flat_direct[idx]->getOutputDimension());
-    EXPECT_EQ(flat_backbone[idx]->getActivationType(),
-              flat_direct[idx]->getActivationType());
-    EXPECT_EQ(flat_backbone[idx]->getName(), flat_direct[idx]->getName());
+    auto &backbone_layer = flat_backbone[idx]->getObject();
+    auto &direct_layer = flat_direct[idx]->getObject();
+    EXPECT_EQ(backbone_layer->getType(), direct_layer->getType());
+    EXPECT_EQ(backbone_layer->getInputDimension(),
+              direct_layer->getInputDimension());
+    EXPECT_EQ(backbone_layer->getOutputDimension(),
+              direct_layer->getOutputDimension());
+    EXPECT_EQ(backbone_layer->getActivationType(),
+              direct_layer->getActivationType());
+    EXPECT_EQ(backbone_layer->getName(), direct_layer->getName());
   }
 }
 
diff --git a/test/unittest/unittest_nntrainer_models.cpp b/test/unittest/unittest_nntrainer_models.cpp
index d57db6f7..ec95caf9 100644
--- a/test/unittest/unittest_nntrainer_models.cpp
+++ b/test/unittest/unittest_nntrainer_models.cpp
@@ -44,7 +44,7 @@ static const std::string getModelsPath(const std::string &file_name) {
  * Watcher Classes                                      *
  ********************************************************/
 
-using NodeType = nntrainer::LayerNode;
+using NodeType = std::shared_ptr<nntrainer::LayerNode>;
 using FlatGraphType = nntrainer::NeuralNetwork::FlatGraphType;
 using NetworkGraphType = nntrainer::NetworkGraph;
 
@@ -91,7 +91,7 @@ public:
    * @brief Construct a new Node Watcher object
    *
    */
-  NodeWatcher() : node(nullptr, 0) {}
+  NodeWatcher() : node(nullptr) {}
 
   /**
    * @brief Construct a new Node Watcher object
@@ -99,21 +99,21 @@ public:
    * @param node node to watch.
    */
   NodeWatcher(const NodeType &node) : node(node) {
-    unsigned int num_weights = node.getObject()->getNumWeights();
-    if (node.getObject()->getType() != nntrainer::InputLayer::type &&
-        node.getObject()->getType() != nntrainer::PreprocessFlipLayer::type &&
-        node.getObject()->getType() !=
+    unsigned int num_weights = node->getObject()->getNumWeights();
+    if (node->getObject()->getType() != nntrainer::InputLayer::type &&
+        node->getObject()->getType() != nntrainer::PreprocessFlipLayer::type &&
+        node->getObject()->getType() !=
           nntrainer::PreprocessTranslateLayer::type)
-      node.getObject()->setTrainable(true);
+      node->getObject()->setTrainable(true);
 
     for (unsigned int i = 0; i < num_weights; ++i) {
-      const nntrainer::Weight &w = node.getObject()->weightAt(i);
+      const nntrainer::Weight &w = node->getObject()->weightAt(i);
       expected_weights.push_back(w.clone());
     }
 
     expected_output =
-      nntrainer::Tensor(node.getObject()->getOutputDimension()[0]);
-    expected_dx = nntrainer::Tensor(node.getObject()->getInputDimension()[0]);
+      nntrainer::Tensor(node->getObject()->getOutputDimension()[0]);
+    expected_dx = nntrainer::Tensor(node->getObject()->getInputDimension()[0]);
   }
 
   /**
@@ -121,9 +121,9 @@ public:
    *
    */
   void readLayerWeight(std::ifstream &f) {
-    for (unsigned int i = 0; i < node.getObject()->getNumWeights(); ++i) {
+    for (unsigned int i = 0; i < node->getObject()->getNumWeights(); ++i) {
       /// @note below is harrasing the fact the tensor shares same base memory
-      node.getObject()->weightAt(i).getVariable().read(f);
+      node->getObject()->weightAt(i).getVariable().read(f);
     }
   }
 
@@ -179,7 +179,7 @@ public:
    *
    * @return float loss
    */
-  float getLoss() { return node.getObject()->getLoss(); }
+  float getLoss() { return node->getObject()->getLoss(); }
 
   /**
    * @brief read Node
@@ -193,10 +193,16 @@ public:
    *
    * @return LayerType
    */
-  std::string getNodeType() { return node.getObject()->getType(); }
+  std::string getNodeType() { return node->getObject()->getType(); }
 
+  /**
+   * @brief get Time Distribution internal layer tyoe
+   *
+   * @return LayerType
+   */
   std::string getTimeDistInternalLayerType() {
-    return std::dynamic_pointer_cast<nntrainer::TimeDistLayer>(node.getObject())
+    return std::dynamic_pointer_cast<nntrainer::TimeDistLayer>(
+             node->getObject())
       ->getDistLayerType();
   }
 
@@ -276,16 +282,16 @@ void NodeWatcher::read(std::ifstream &in) {
 
 void NodeWatcher::verifyWeight(const std::string &error_msg) {
   for (unsigned int i = 0; i < expected_weights.size(); ++i) {
-    verify(node.getObject()->weightAt(i).getVariable(),
+    verify(node->getObject()->weightAt(i).getVariable(),
            expected_weights[i].getVariable(),
-           error_msg + " " + node.getObject()->weightAt(i).getName() +
+           error_msg + " " + node->getObject()->weightAt(i).getName() +
              " weight");
   }
 }
 
 void NodeWatcher::verifyGrad(const std::string &error_msg) {
   for (unsigned int i = 0; i < expected_weights.size(); ++i) {
-    auto weight = node.getObject()->weightAt(i);
+    auto weight = node->getObject()->weightAt(i);
     if (weight.getTrainable()) {
       verify(weight.getGradient(), expected_weights[i].getGradient(),
              error_msg + " " + weight.getName() + " grad");
@@ -295,19 +301,19 @@ void NodeWatcher::verifyGrad(const std::string &error_msg) {
 
 void NodeWatcher::forward(int iteration, NodeWatcher &next_node) {
   std::stringstream ss;
-  ss << "forward failed at " << node.getObject()->getName() << " at iteration "
+  ss << "forward failed at " << node->getObject()->getName() << " at iteration "
      << iteration;
   std::string err_msg = ss.str();
 
-  std::vector<nntrainer::Tensor> out = node.getObject()->getOutputs();
+  std::vector<nntrainer::Tensor> out = node->getObject()->getOutputs();
 
   /**
    * @todo Do not veify if the layer is operting in-place by checking its
    * property
    */
-  if (next_node.node.getObject()->getType() !=
+  if (next_node.node->getObject()->getType() !=
         nntrainer::ActivationLayer::type &&
-      next_node.node.getObject()->getType() !=
+      next_node.node->getObject()->getType() !=
         nntrainer::BatchNormalizationLayer::type)
     verify(out[0], expected_output, err_msg + " at output");
 }
@@ -316,12 +322,12 @@ nntrainer::sharedConstTensors
 NodeWatcher::lossForward(nntrainer::sharedConstTensors pred,
                          nntrainer::sharedConstTensors answer, int iteration) {
   std::stringstream ss;
-  ss << "loss failed at " << node.getObject()->getName() << " at iteration "
+  ss << "loss failed at " << node->getObject()->getName() << " at iteration "
      << iteration;
   std::string err_msg = ss.str();
 
   nntrainer::sharedConstTensors out =
-    std::static_pointer_cast<nntrainer::LossLayer>(node.getObject())
+    std::static_pointer_cast<nntrainer::LossLayer>(node->getObject())
       ->forwarding_with_val(pred, answer);
 
   return out;
@@ -329,11 +335,11 @@ NodeWatcher::lossForward(nntrainer::sharedConstTensors pred,
 
 void NodeWatcher::backward(int iteration, bool verify_deriv, bool verify_grad) {
   std::stringstream ss;
-  ss << "backward failed at " << node.getObject()->getName() << " at iteration "
-     << iteration;
+  ss << "backward failed at " << node->getObject()->getName()
+     << " at iteration " << iteration;
   std::string err_msg = ss.str();
 
-  std::vector<nntrainer::Tensor> out = node.getObject()->getDerivatives();
+  std::vector<nntrainer::Tensor> out = node->getObject()->getDerivatives();
 
   if (verify_grad) {
     verifyGrad(err_msg + " grad");
@@ -379,7 +385,7 @@ GraphWatcher::GraphWatcher(const std::string &config, const bool opt) :
   std::vector<NodeType> graph = model_graph.getSorted();
 
   for (auto it = graph.begin(); it != graph.end() - 1; ++it) {
-    nodes.push_back(NodeWatcher(*it));
+    nodes.push_back(NodeWatcher((*it)));
   }
 
   loss_node = NodeWatcher(graph.back());
@@ -1059,8 +1065,10 @@ TEST(nntrainerModels, read_save_01_n) {
     nntrainer::createLayer(nntrainer::InputLayer::type);
   layer->setProperty(
     {"input_shape=1:1:62720", "normalization=true", "bias_initializer=zeros"});
+  std::shared_ptr<nntrainer::LayerNode> layer_node =
+    std::make_unique<nntrainer::LayerNode>(layer);
 
-  EXPECT_NO_THROW(NN.addLayer(layer));
+  EXPECT_NO_THROW(NN.addLayer(layer_node));
   EXPECT_NO_THROW(NN.setProperty({"loss=mse"}));
 
   EXPECT_THROW(NN.readModel(), std::runtime_error);