[onert_train] Introduce some arguments for checkpoint apis (#13744)

This commit introduces `--checkpoint` argument for loading a checkpoint file and `--export_checkpoint` argument for exporting weight and bias data of the model to checkpoint file. ONE-DCO-1.0-Signed-off-by: Jiyoung Yun <[email protected]>
Samsung · Aug 26, 2024 · 38d61b4 · 38d61b4
1 parent abcfcb5
commit 38d61b4
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/tests/tools/onert_train/src/args.cc b/tests/tools/onert_train/src/args.cc
@@ -146,10 +146,14 @@ void Args::Initialize(void)
     .nargs(0)
     .default_value(false)
     .help("Print version and exit immediately");
+  _arser.add_argument("--checkpoint").type(arser::DataType::STR).help("Checkpoint filename");
   _arser.add_argument("--export_circle").type(arser::DataType::STR).help("Path to export circle");
   _arser.add_argument("--export_circleplus")
     .type(arser::DataType::STR)
     .help("Path to export circle+");
+  _arser.add_argument("--export_checkpoint")
+    .type(arser::DataType::STR)
+    .help("Path to export checkpoint");
   _arser.add_argument("--load_input:raw")
     .type(arser::DataType::STR)
     .help({"NN Model Raw Input data file", "The datafile must have data for each input number.",
@@ -243,10 +247,18 @@ void Args::Parse(const int argc, char **argv)
       }
     }
 
+    if (_arser["--checkpoint"])
+    {
+      _checkpoint_filename = _arser.get<std::string>("--checkpoint");
+      checkModelfile(_checkpoint_filename);
+    }
+
     if (_arser["--export_circle"])
       _export_circle_filename = _arser.get<std::string>("--export_circle");
     if (_arser["--export_circleplus"])
       _export_circleplus_filename = _arser.get<std::string>("--export_circleplus");
+    if (_arser["--export_checkpoint"])
+      _export_checkpoint_filename = _arser.get<std::string>("--export_checkpoint");
     if (_arser["--load_input:raw"])
     {
       _load_raw_input_filename = _arser.get<std::string>("--load_input:raw");

diff --git a/tests/tools/onert_train/src/args.h b/tests/tools/onert_train/src/args.h
@@ -49,8 +49,10 @@ class Args
 
   const std::string &getPackageFilename(void) const { return _package_filename; }
   const std::string &getModelFilename(void) const { return _model_filename; }
+  const std::string &getCheckpointFilename(void) const { return _checkpoint_filename; }
   const std::string &getExportCircleFilename(void) const { return _export_circle_filename; }
   const std::string &getExportCirclePlusFilename(void) const { return _export_circleplus_filename; }
+  const std::string &getExportCheckpointFilename(void) const { return _export_checkpoint_filename; }
   const bool useSingleModel(void) const { return _use_single_model; }
   const std::string &getLoadRawInputFilename(void) const { return _load_raw_input_filename; }
   const std::string &getLoadRawExpectedFilename(void) const { return _load_raw_expected_filename; }
@@ -99,8 +101,10 @@ class Args
 
   std::string _package_filename;
   std::string _model_filename;
+  std::string _checkpoint_filename;
   std::string _export_circle_filename;
   std::string _export_circleplus_filename;
+  std::string _export_checkpoint_filename;
   bool _use_single_model = false;
   std::string _load_raw_input_filename;
   std::string _load_raw_expected_filename;

diff --git a/tests/tools/onert_train/src/onert_train.cc b/tests/tools/onert_train/src/onert_train.cc
@@ -156,7 +156,12 @@ int main(const int argc, char **argv)
     // prepare execution
 
     // TODO When nnfw_{prepare|run} are failed, can't catch the time
-    measure.run(PhaseType::PREPARE, [&]() { NNPR_ENSURE_STATUS(nnfw_train_prepare(session)); });
+    measure.run(PhaseType::PREPARE, [&]() {
+      NNPR_ENSURE_STATUS(nnfw_train_prepare(session));
+
+      if (const auto name = args.getCheckpointFilename(); name != "")
+        NNPR_ENSURE_STATUS(nnfw_train_import_checkpoint(session, name.c_str()));
+    });
 
     // prepare input and expected tensor info lists
     std::vector<nnfw_tensorinfo> input_infos;
@@ -280,6 +285,9 @@ int main(const int argc, char **argv)
               if (args.getMetricType() == 0)
                 metrics[i] += metric.categoricalAccuracy(i);
             }
+
+            if (const auto name = args.getExportCheckpointFilename(); name != "")
+              NNPR_ENSURE_STATUS(nnfw_train_export_checkpoint(session, name.c_str()));
           }
 
           // print loss