From 30ec54e6c7a58feb292471c576ef7b61aa718c70 Mon Sep 17 00:00:00 2001
From: Parichay Kapoor <pk.kapoor@samsung.com>
Date: Wed, 3 Mar 2021 20:28:17 +0900
Subject: [PATCH] [neuralnet] Support deallocate of tensors

Support deallocate and allocation of tensors from the neuralnet.

Also perform the deallocation of tensors after each train run.
Once a training is performed, the memory associated for that training
except the weights variables will be freed.

The memory associated with inference will not be freed until freed manually.
This will require calling deallocate() with the model object.
Note that calling a train() after inference() will free the inference memory
and the train will allocate its own memory which will be freed at the end
of the training.

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

Signed-off-by: Parichay Kapoor <pk.kapoor@samsung.com>
---
 nntrainer/models/neuralnet.cpp              | 22 ++++++++++++++++++----
 nntrainer/models/neuralnet.h                | 13 ++++++++++---
 test/unittest/unittest_nntrainer_graph.cpp  |  2 +-
 test/unittest/unittest_nntrainer_models.cpp |  2 +-
 4 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/nntrainer/models/neuralnet.cpp b/nntrainer/models/neuralnet.cpp
index 4a724ac..6b22d13 100644
--- a/nntrainer/models/neuralnet.cpp
+++ b/nntrainer/models/neuralnet.cpp
@@ -527,7 +527,7 @@ sharedConstTensors NeuralNetwork::inference(sharedConstTensors X) {
   if (!validateInput(X))
     return out;
 
-  assignMem(false);
+  allocate(false);
 
   try {
     START_PROFILE(profile::NN_FORWARD);
@@ -552,7 +552,7 @@ sharedConstTensors NeuralNetwork::inference(sharedConstTensors X) {
   return out;
 }
 
-int NeuralNetwork::assignMem(bool trainable) {
+int NeuralNetwork::allocate(bool trainable) {
   // TODO: directly replace this
   manager->initializeTensors(trainable);
 
@@ -560,6 +560,12 @@ int NeuralNetwork::assignMem(bool trainable) {
   return ML_ERROR_NONE;
 }
 
+int NeuralNetwork::deallocate() {
+  manager->deallocateTensors(true);
+
+  return ML_ERROR_NONE;
+}
+
 int NeuralNetwork::train(std::vector<std::string> values) {
   int status = ML_ERROR_NONE;
 
@@ -574,7 +580,7 @@ int NeuralNetwork::train(std::vector<std::string> values) {
   /** set batch size just before training */
   setBatchSize(batch_size);
 
-  status = assignMem(true);
+  status = allocate(true);
   NN_RETURN_STATUS();
 
   /** Setup data buffer properties */
@@ -587,7 +593,15 @@ int NeuralNetwork::train(std::vector<std::string> values) {
   status = data_buffer->init();
   NN_RETURN_STATUS();
 
-  return train_run();
+  status = train_run();
+
+  /**
+   * Free the memory needed for training before exiting.
+   * Note that this does not free the weights for the model.
+   * Weights of the model will be freed when the model is destroyed.
+   */
+  manager->deallocateTensors(false);
+  return status;
 }
 
 /**
diff --git a/nntrainer/models/neuralnet.h b/nntrainer/models/neuralnet.h
index c7784ab..20c49cb 100644
--- a/nntrainer/models/neuralnet.h
+++ b/nntrainer/models/neuralnet.h
@@ -162,13 +162,20 @@ public:
   int initialize();
 
   /**
-   * @brief     Assign Graph Memory. This should be called after initialize.
-   * @TODO      Consider to move Network Graph Class
+   * @brief     Allocate memory for the model. This should be called after initialize.
    * @param[in] trainable Assign memory for inference or train mode
    * @retval #ML_ERROR_NONE Successful.
    * @retval #ML_ERROR_INVALID_PARAMETER invalid parameter.
    */
-  int assignMem(bool trainable = true);
+  int allocate(bool trainable = true);
+
+  /**
+   * @brief     Deallocate memory for the model.
+   * @param[in] trainable Assign memory for inference or train mode
+   * @retval #ML_ERROR_NONE Successful.
+   * @note This does not free the model graph but only the weight tensors, and input/output/gradient/derivative tensors if any.
+   */
+  int deallocate();
 
   /**
    * @brief     Update graph to make batch normalization in-place
diff --git a/test/unittest/unittest_nntrainer_graph.cpp b/test/unittest/unittest_nntrainer_graph.cpp
index 33f70e4..7ab32a8 100644
--- a/test/unittest/unittest_nntrainer_graph.cpp
+++ b/test/unittest/unittest_nntrainer_graph.cpp
@@ -88,7 +88,7 @@ TEST_P(nntrainerGraphTest, loadConfig) {
     EXPECT_EQ(status, ML_ERROR_NONE);
   }
 
-  status = NN.assignMem();
+  status = NN.allocate();
   if (failAtLoad()) {
     EXPECT_NE(status, ML_ERROR_NONE);
   } else {
diff --git a/test/unittest/unittest_nntrainer_models.cpp b/test/unittest/unittest_nntrainer_models.cpp
index 63f0824..b267add 100644
--- a/test/unittest/unittest_nntrainer_models.cpp
+++ b/test/unittest/unittest_nntrainer_models.cpp
@@ -322,7 +322,7 @@ GraphWatcher::GraphWatcher(const std::string &config, const bool opt) :
     throw std::invalid_argument("initiation failed");
   };
 
-  if (nn.assignMem()) {
+  if (nn.allocate()) {
     throw std::invalid_argument("assign Memory failed");
   };
 
-- 
2.7.4