Merge pull request brucefan1983#1142 from brucefan1983/dump-nep-restart-periodically

Dump nep restart periodically

Author: brucefan1983

doc/nep/input_parameters/save_potential.rst

@@ -9,10 +9,11 @@ This keyword sets the number of of generations between writing a ``nep.txt`` che
 If ``<format>`` is set to ``0`` the output file name is formatted as ``nep_gen[generation].txt``.
 If ``<format>`` is set to ``1`` the output file name is formatted as ``nep_y[year]_m[month]_d[day]_h[hour]_m[minute]_s[second]_generation[generation].txt``.
 These model files can be used to monitor the training progress of your model.
+Additionally, if ``<save_restart>`` is set to ``1`` the :ref:`nep.restart file <nep_restart>` is also written, following the naming format set by ``<format>``.
 Note that the :ref:`nep.restart file <nep_restart>` is the file that is required to continue training.
 
 The syntax is::
 
-  save_potential <number_of_generations_between_save_potential> <format>
+  save_potential <number_of_generations_between_save_potential> <format> <save_restart>
 
 The default number of generations between saved model files is :math:`N=10^5` and the output file names use the extended format.

src/main_nep/fitness.cu

@@ -415,6 +415,19 @@ void Fitness::write_nep_txt(FILE* fid_nep, Parameters& para, float* elite)
   }
 }
 
+void Fitness::get_save_potential_label(Parameters& para, const int generation, std::string& label) {
+    if (para.save_potential_format == 1) {
+      time_t rawtime;
+      time(&rawtime);
+      struct tm* timeinfo = localtime(&rawtime);
+      char buffer[200];
+      strftime(buffer, sizeof(buffer), "nep_y%Y_m%m_d%d_h%H_m%M_s%S_generation", timeinfo);
+      label = std::string(buffer) + std::to_string(generation + 1);
+    } else {
+      label = "nep_gen" + std::to_string(generation + 1);
+    }
+}
+
 void Fitness::report_error(
   Parameters& para,
   const int generation,
@@ -462,16 +475,8 @@ void Fitness::report_error(
 
     if (0 == (generation + 1) % para.save_potential) {
       std::string filename;
-      if (para.save_potential_format == 1) {
-        time_t rawtime;
-        time(&rawtime);
-        struct tm* timeinfo = localtime(&rawtime);
-        char buffer[200];
-        strftime(buffer, sizeof(buffer), "nep_y%Y_m%m_d%d_h%H_m%M_s%S_generation", timeinfo);
-        filename = std::string(buffer) + std::to_string(generation + 1) + ".txt";
-      } else {
-        filename = "nep_gen" + std::to_string(generation + 1) + ".txt";
-      }
+      get_save_potential_label(para, generation, filename);
+      filename += ".txt";
 
       FILE* fid_nep = my_fopen(filename.c_str(), "w");
       write_nep_txt(fid_nep, para, elite);

src/main_nep/fitness.cuh

@@ -37,6 +37,7 @@ public:
     const float loss_L2,
     float* elite);
   void predict(Parameters& para, float* elite);
+  void get_save_potential_label(Parameters& para, const int generation, std::string& filename);
 
 protected:
   bool has_test_set = false;

src/main_nep/parameters.cu

@@ -1316,8 +1316,8 @@ void Parameters::parse_save_potential(const char** param, int num_param)
 {
   is_save_potential_set = true;
 
-  if (num_param != 3) {
-    PRINT_INPUT_ERROR("save_potential should have 2 parameters.\n");
+  if (num_param != 4) {
+    PRINT_INPUT_ERROR("save_potential should have 3 parameters.\n");
   }
   if (!is_valid_int(param[1], &save_potential)) {
     PRINT_INPUT_ERROR("save_potential interval should be an integer.\n");
@@ -1331,4 +1331,10 @@ void Parameters::parse_save_potential(const char** param, int num_param)
   if (save_potential_format != 0 && save_potential_format != 1) {
     PRINT_INPUT_ERROR("save_potential format should be 0 or 1.");
   }  
+  if (!is_valid_int(param[3], &save_potential_restart)) {
+    PRINT_INPUT_ERROR("save_potential save restart should be an integer.\n");
+  }
+  if (save_potential_restart != 0 && save_potential_restart != 1) {
+    PRINT_INPUT_ERROR("save_potential save restart should be 0 or 1.");
+  }  
 }

src/main_nep/parameters.cuh

@@ -31,7 +31,8 @@ public:
   int population_size;    // population size for SNES
   int maximum_generation; // maximum number of generations for SNES;
   int save_potential;     // number of generations between writing a checkpoint nep.txt file. 
-  int save_potential_format;  // format of checkpoint nep.txt file name
+  int save_potential_format;   // format of checkpoint nep.txt file name
+  int save_potential_restart;  // if restart files should be written or not. 0=no, 1=yes
   int num_neurons1;       // number of nuerons in the 1st hidden layer (only one hidden layer)
   int basis_size_radial;  // for nep3
   int basis_size_angular; // for nep3

src/main_nep/snes.cu

@@ -356,7 +356,15 @@ void SNES::compute(Parameters& para, Fitness* fitness_function)
 
       update_mu_and_sigma(para);
       if (0 == (n + 1) % 100) {
-        output_mu_and_sigma(para);
+        const char* filename = "nep.restart";
+        output_mu_and_sigma(para, filename);
+      }
+      // Optionally save the nep.restart file at the same time as save_potential
+      if (0 == (n + 1) % para.save_potential && para.save_potential_restart) {
+        std::string restart_file;
+        fitness_function->get_save_potential_label(para, n, restart_file);
+        restart_file += ".restart";
+        output_mu_and_sigma(para, restart_file.c_str());
       }
     }
   } else {
@@ -642,12 +650,12 @@ void SNES::update_mu_and_sigma(Parameters& para)
   GPU_CHECK_KERNEL;
 }
 
-void SNES::output_mu_and_sigma(Parameters& para)
+void SNES::output_mu_and_sigma(Parameters& para, const char* filename)
 {
   gpuSetDevice(0); // normally use GPU-0
   gpu_mu.copy_to_host(mu.data());
   gpu_sigma.copy_to_host(sigma.data());
-  FILE* fid_restart = my_fopen("nep.restart", "w");
+  FILE* fid_restart = my_fopen(filename, "w");
   for (int n = 0; n < number_of_variables; ++n) {
     fprintf(fid_restart, "%15.7e %15.7e\n", mu[n], sigma[n]);
   }

src/main_nep/snes.cuh

@@ -70,5 +70,5 @@ protected:
   void regularize_NEP4(Parameters& para);
   void sort_population(Parameters& para);
   void update_mu_and_sigma(Parameters& para);
-  void output_mu_and_sigma(Parameters& para);
+  void output_mu_and_sigma(Parameters& para, const char* filename);
 };