[onert-micro] Introduce TrainingModule class

onert-micro/luci-interpreter/src/core/CMakeLists.txt

@@ -6,7 +6,8 @@ set(SOURCES
     RuntimeModule.h
     TrainingGraph.cpp
     GradientCalculationStorage.cpp
-    TrainableWeightStorage.cpp)
+    TrainableWeightStorage.cpp
+    TrainingModule.cpp)
 
 add_library(${LUCI_INTERPRETER_CORE} STATIC ${SOURCES})
 if (NOT NNCC_LIBRARY_NO_PIC)

onert-micro/luci-interpreter/src/core/TrainingModule.cpp

@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifdef ENABLE_TRAINING
+
+#include "TrainingModule.h"
+
+#include <memory>
+
+namespace luci_interpreter
+{
+namespace training
+{
+
+training::Status TrainingModule::enableTrainingMode(training::TrainingSettings &settings,
+                                                    SimpleMemoryManager *memoryManager)
+{
+  if (_runtime_module->_storage.get() == nullptr)
+  {
+    _runtime_module->_storage = std::make_unique<TrainableWeightStorage>();
+  }
+
+  if (_runtime_module->_storage->fillTrainableWeightsStorage(
+        &_runtime_module->_circle_reader, memoryManager,
+        settings.number_of_last_trainable_layers) == training::Error)
+    return training::Error;
+
+  _training_graph = std::make_unique<training::TrainingGraph>();
+
+  for (auto &graph : _runtime_module->_graphs)
+  {
+    graph.setLastTrainingLayersNumber(settings.number_of_last_trainable_layers);
+    graph.setGradientCalculationStorage(_training_graph->getGradientCalculationStorage());
+    graph.setTrainingWeightStorage(_runtime_module->_storage.get());
+  }
+
+  return training::Ok;
+}
+
+training::Status TrainingModule::disableTrainingMode(bool resetWeights)
+{
+  _training_graph.release();
+
+  if (resetWeights)
+  {
+    if (_runtime_module->_storage->clearAllTrainableWeights() == training::Error)
+      return training::Error;
+    _runtime_module->_storage.release();
+  }
+
+  for (auto &graph : _runtime_module->_graphs)
+  {
+    graph.setLastTrainingLayersNumber(0);
+    graph.setGradientCalculationStorage(nullptr);
+    if (resetWeights)
+      graph.setTrainingWeightStorage(nullptr);
+  }
+
+  return training::Ok;
+}
+
+} // namespace training
+} // namespace luci_interpreter
+
+#endif // ENABLE_TRAINING

onert-micro/luci-interpreter/src/core/TrainingModule.h

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 2023 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifdef ENABLE_TRAINING
+
+#ifndef LUCI_INTERPRETER_CORE_TRAINING_MODULE_H
+#define LUCI_INTERPRETER_CORE_TRAINING_MODULE_H
+
+#include "core/RuntimeModule.h"
+
+#include "luci_interpreter/core/TrainableWeightStorage.h"
+#include "TrainingGraph.h"
+
+namespace luci_interpreter
+{
+namespace training
+{
+
+class TrainingModule
+{
+public:
+  TrainingModule(RuntimeModule *runtime_module) : _runtime_module(runtime_module)
+  {
+    // Do nothing
+  }
+
+  training::Status enableTrainingMode(training::TrainingSettings &settings,
+                                      SimpleMemoryManager *memoryManager);
+
+  training::Status disableTrainingMode(bool resetWeights);
+
+  training::Status computeGradients(const TrainingSettings &settings,
+                                    const uint8_t *label_train_data)
+  {
+    return _training_graph->computeGradients(settings, _runtime_module->_storage.get(),
+                                             &_runtime_module->_circle_reader, label_train_data);
+  }
+
+  training::Status updateWeights(const TrainingSettings &settings)
+  {
+    return _training_graph->updateWeights(settings, _runtime_module->_storage.get(),
+                                          &_runtime_module->_circle_reader);
+  }
+
+private:
+  RuntimeModule *_runtime_module;
+
+  std::unique_ptr<training::TrainingGraph> _training_graph;
+};
+
+} // namespace training
+} // namespace luci_interpreter
+
+#endif // LUCI_INTERPRETER_CORE_TRAINING_MODULE_H
+
+#endif // ENABLE_TRAINING