[networkgraph] Network graph updated for Layer V2

Network graph updated to work with LayerV2
This involves settings input dimension in InitContext, and setting
up RunContext for each layer before their execution.

Further, some helper functions are also added in LayerNode.

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>

diff --git a/nntrainer/graph/network_graph.cpp b/nntrainer/graph/network_graph.cpp
index d02b74d..63e81b6 100644 (file)
--- a/nntrainer/graph/network_graph.cpp
+++ b/nntrainer/graph/network_graph.cpp
@@ -379,7 +379,7 @@ int NetworkGraph::checkCompiledGraph() {
   /** Dimension of input layers must be known */
   for (auto iter = cbegin(); iter != cend(); iter++) {
     auto lnode = (*iter);
-    if (lnode->getObject()->getType() == InputLayer::type) {
+    if (lnode->getType() == InputLayer::type) {
       if (lnode->getInputDimensions().size() == 0) {
         ml_loge("InputDimension of first layer is not set");
         return ML_ERROR_INVALID_PARAMETER;
@@ -696,45 +696,64 @@ void NetworkGraph::inPlaceOptimize(Manager &manager) {
   }
 }
 
-void NetworkGraph::updateRunContext(std::shared_ptr<Manager> &manager,
-                                    const std::shared_ptr<LayerNode> &lnode) {
+std::vector<Var_Grad *>
+NetworkGraph::updateRunContext(std::shared_ptr<Manager> &manager,
+                               const std::shared_ptr<LayerNode> &lnode,
+                               const std::vector<Var_Grad *> &prev_inputs) {
   /**
    * using copy assignment allows setting run_context without adding more
    * interfaces
    */
   const GraphNode &gnode = *lnode.get();
   const InitLayerContext &init_context = lnode->getInitContext();
+  std::vector<Var_Grad *> inputs = prev_inputs;
+  if (inputs.empty())
+    inputs = manager->requestInputs(gnode, init_context.getInputDimensions());
+
+  const std::vector<Var_Grad *> &outputs =
+    manager->requestOutputs(gnode, init_context.getOutputDimensions());
+
   /**
    * @todo must use existing properties like name/trainable of run_context to
    * create the new run_context
    */
   // const RunLayerContext &run_context = lnode->getRunContext();
-
   lnode->updateRunContext(RunLayerContext(
-    manager->requestWeights(gnode, init_context.getWeightsSpec()),
-    manager->requestInputs(gnode, init_context.getInputDimensions()),
-    manager->requestOutputs(gnode, init_context.getOutputDimensions()),
-    manager->requestTensors(gnode, init_context.getTensorsSpec())));
+    manager->requestWeights(gnode, init_context.getWeightsSpec()), inputs,
+    outputs, manager->requestTensors(gnode, init_context.getTensorsSpec())));
+
+  return outputs;
 }
 
 int NetworkGraph::initialize(std::shared_ptr<Manager> manager) {
   int status = ML_ERROR_NONE;
-  // TODO: don't delete adj list - use it here and make this more cleaner/faster
+  /** this contains the map from name to input tensors for each node */
+  std::unordered_map<std::string, std::vector<Var_Grad *>> input_map;
+
+  /** check if the given config of node is of input node */
+  auto is_input_node = [](const std::string &type,
+                          const unsigned int idx) -> bool {
+    /** TODO: remove dependency on idx */
+    return type == InputLayer::type || idx == 0;
+  };
 
   for (unsigned int idx = 0; idx < graph.size(); ++idx) {
-    bool first = idx == 0;
     auto const &lnode = getSortedLayerNode(idx);
-    auto &lptr = lnode->getObject();
+    std::string cur_type = lnode->getType();
     ml_logd("layer name : %s", lnode->getName().c_str());
-    std::string cur_type = lptr->getType();
+
+#if !LAYER_V2
+    auto &lptr = lnode->getObject();
+#endif
 
     /**
      * Set input dimension for all the layers.
      * For input layer, as input dimension is known, set input tensor.
      */
-    if (!first) {
+    if (!is_input_node(cur_type, idx)) {
       std::string l_pre_type = getSortedLayerNode(idx - 1)->getType();
 
+      // TODO: move this to checkCompiledGraph
       if (istrequal(l_pre_type, ActivationLayer::type) &&
           istrequal(cur_type, ActivationLayer::type)) {
         ml_loge("double activation is not allowed");
@@ -753,9 +772,13 @@ int NetworkGraph::initialize(std::shared_ptr<Manager> manager) {
           }
         }
 
-        // TODO: set init layer context init layer
+#if LAYER_V2
+        lnode->setInputDimension(in_layer_node->getOutputDimensions()[location],
+                                 i);
+#else
         lptr->setInputDimension(in_layer_node->getOutputDimensions()[location],
                                 i);
+#endif
       }
     }
 
@@ -763,23 +786,53 @@ int NetworkGraph::initialize(std::shared_ptr<Manager> manager) {
      * Initialize all the layers, allocate output tensors for each layer
      * init2and add optimizer related weights for the layer
      */
-    // TODO: this call will fill the init context inside the layer
-    // lnode->initialize();
+#if LAYER_V2
+    lnode->finalize();
+#else
     status = lptr->initialize(*manager);
     NN_RETURN_STATUS();
-
-    updateRunContext(manager, lnode);
-    // TODO: remove this
+#endif
+
+#if LAYER_V2
+    std::vector<Var_Grad *> inputs = {};
+    if (!is_input_node(cur_type, idx)) {
+      if (input_map.find(lnode->getName()) == input_map.end())
+        throw std::runtime_error("Cannot find input buffers for the node");
+      inputs = input_map.at(lnode->getName());
+    }
+    const std::vector<Var_Grad *> &outputs =
+      updateRunContext(manager, lnode, inputs);
+#else
     auto &in_out = manager->trackLayerOutputs(cur_type, lnode->getName(),
                                               lptr->getOutputDimension(),
                                               lptr->getInputDimension());
-    // TODO: remove this
     lptr->setOutputBuffers(in_out);
+#endif
+
+#if LAYER_V2
+    auto &output_layers = lnode->getOutputLayers();
+    for (unsigned int i = 0; i < outputs.size(); ++i) {
+      auto out_layer_node = getLayerNode(output_layers[i]);
+      if (input_map.find(output_layers[i]) == input_map.end())
+        input_map.insert({output_layers[i], {}});
+
+      unsigned int j = 0;
+      for (; j < out_layer_node->getNumInputs(); ++j) {
+        if (out_layer_node->getInputLayers()[j] == lnode->getName()) {
+          break;
+        }
+      }
 
-    // TODO: fill runcontext of other guys as well
-    /** Connect the output of the previous layers with the input of the current
-     * layer */
-    if (!first) {
+      auto &in_map = input_map.at(output_layers[i]);
+      in_map.resize(out_layer_node->getNumInputs());
+      in_map[j] = outputs[i];
+    }
+#else
+    /**
+     * Connect the output of the previous layers with the input of the current
+     * layer
+     */
+    if (!is_input_node(cur_type, idx)) {
       auto &input_layers = lnode->getInputLayers();
       for (unsigned int i = 0; i < input_layers.size(); ++i) {
         auto in_layer_node = getLayerNode(input_layers[i]);
@@ -804,6 +857,7 @@ int NetworkGraph::initialize(std::shared_ptr<Manager> manager) {
                                                lptr->getOutputDimension());
       lptr->setInputBuffers(in_out);
     }
+#endif
   }
   return status;
 }

diff --git a/nntrainer/graph/network_graph.h b/nntrainer/graph/network_graph.h
index 9573ba4..9ab3292 100644 (file)
--- a/nntrainer/graph/network_graph.h
+++ b/nntrainer/graph/network_graph.h
@@ -262,8 +262,10 @@ public:
    * @param init_context Init layer context to create run context
    * @param run_context Run layer context to be created
    */
-  static void updateRunContext(std::shared_ptr<Manager> &manager,
-                               const std::shared_ptr<LayerNode> &lnode);
+  static std::vector<Var_Grad *>
+  updateRunContext(std::shared_ptr<Manager> &manager,
+                   const std::shared_ptr<LayerNode> &lnode,
+                   const std::vector<Var_Grad *> &inputs);
 
 private:
   std::map<std::string, std::string> sub_in_out; /** This is map to identify

diff --git a/nntrainer/layers/layer_context.h b/nntrainer/layers/layer_context.h
index d466993..c43bf53 100644 (file)
--- a/nntrainer/layers/layer_context.h
+++ b/nntrainer/layers/layer_context.h
@@ -61,6 +61,13 @@ public:
   InitLayerContext() = default;
 
   /**
+   * @brief Construct a new Init Layer Context object
+   *
+   * @param dim Input dimensions for the layer
+   */
+  InitLayerContext(const std::vector<TensorDim> &dim) : input_dim(dim) {}
+
+  /**
    * @brief Get the Input Dimensions object
    *
    * @return const std::vector<TensorDim>& Input dimensions

diff --git a/nntrainer/layers/layer_internal.h b/nntrainer/layers/layer_internal.h
index 07e2749..5bd0afc 100644 (file)
--- a/nntrainer/layers/layer_internal.h
+++ b/nntrainer/layers/layer_internal.h
@@ -578,9 +578,8 @@ public:
    * @param[in] i axis
    */
   void setInputDimension(const TensorDim &d, unsigned int i) {
-    if (i > MAXDIM)
-      throw std::invalid_argument(
-        "axis must be greater than 0 and less then MAX_DIM : 4");
+    if (i >= getNumInputs())
+      throw std::out_of_range("Setting dimensions out of bounds");
     input_dim[i] = d;
   }
 

diff --git a/nntrainer/layers/layer_node.h b/nntrainer/layers/layer_node.h
index 41e80b1..1b6e1b7 100644 (file)
--- a/nntrainer/layers/layer_node.h
+++ b/nntrainer/layers/layer_node.h
@@ -429,6 +429,30 @@ public:
     run_context = std::move(context);
   }
 
+  /**
+   * @brief Set input dimension for the layer
+   *
+   * @param dim Input tensor dim
+   * @param idx Index of the dim
+   */
+  void setInputDimension(const TensorDim &dim, unsigned int idx) {
+    if (idx >= getNumInputs())
+      throw std::out_of_range("Setting dimensions out of bounds");
+    input_dim[idx] = dim;
+  }
+
+  /**
+   * @brief Finalize the layer
+   *
+   */
+  void finalize() {
+    /** Create init context right before finalize */
+    init_context = InitLayerContext(input_dim);
+#if LAYER_V2
+    layer->finalize(init_context);
+#endif
+  }
+
 private:
   // TODO: make this unique_ptr once getObject API is removed
   std::shared_ptr<nntrainer::LayerV1>
@@ -437,7 +461,6 @@ private:
   // can lead to issues later
   size_t index; /**< index of each node */
 
-  /** TODO : move management of num_inputs to layer_node */
   std::vector<std::string> input_layers;  /**< input layer names */
   std::vector<std::string> output_layers; /**< output layer names */
   bool flatten;    /**< flatten the output of this node */
@@ -445,15 +468,19 @@ private:
   ActivationType
     activation_type; /**< activation applied to the output of this node */
 
-  RunLayerContext
-    run_context; /**< context required for running/execution of the layer. This
-                    will also contain the properties of the layer. The
-                    properties will be copied upon final creation. Editing
-                    properties of the layer after init will not the properties
-                    in the context/graph unless intended. */
+  std::vector<TensorDim>
+    input_dim; /**< input dimension for the layer. This can be in partial state
+                  before the layer is initialized */
   InitLayerContext init_context; /**< context to be built for/while
                                     initialization of the layer. This will also
                                     contain the properties of the layer. */
+
+  RunLayerContext run_context; /**< context required for running/execution of
+                    the layer. This will also contain the properties of the
+                    layer. The properties will be copied upon final creation.
+                    Editing properties of the layer after init will not the
+                    properties in the context/graph unless intended. */
+
   /**
    * These properties are set for the layer by the user but are intercepted
    * and used in the node which forms the basic element of the graph.