update model saving

Saurav Agarwal · Saurav Agarwal · commit f4fb93cc58fe · 2025-03-21T13:21:18.000-04:00
diff --git a/python/coverage_control/nn/trainers/trainer.py b/python/coverage_control/nn/trainers/trainer.py
@@ -139,6 +139,7 @@ def train(self) -> None:
                 if val_loss < best_val_loss:
                     best_val_loss = val_loss
                     best_model_state_dict = deepcopy(self.model.state_dict())
+                    best_model_data = {"epoch": epoch, "optimizer_state_dict": deepcopy(self.optimizer.state_dict()), "loss": val_loss}
                     # torch.save(self.model.state_dict(), self.model_dir + "/model.pt")
                     # torch.save(self.optimizer.state_dict(), self.model_dir + "/optimizer.pt")
                 print(f"Epoch: {epoch + 1}/{self.num_epochs} ",
@@ -148,18 +149,18 @@ def train(self) -> None:
             if train_loss < best_train_loss:
                 best_train_loss = train_loss
                 best_train_model_state_dict = deepcopy(self.model.state_dict())
-                # torch.save(self.model.state_dict(), self.model_dir + "/model_curr.pt")
-                # torch.save(self.optimizer.state_dict(), self.model_dir + "/optimizer_curr.pt")
+                best_train_model_data = {"epoch": epoch, "optimizer_state_dict": deepcopy(self.optimizer.state_dict()), "loss": train_loss}
 
-            if epoch % 5 == 0:
+            if (epoch + 1) % 5 == 0:
                 model_state_dict = self.model.state_dict()
-                model_state_dict["epoch"] = epoch
-                model_state_dict["optimizer_state_dict"] = self.optimizer.state_dict()
-                model_state_dict["loss"] = train_loss
                 torch.save(model_state_dict, self.model_dir + "/model_epoch" + str(epoch) + ".pt")
+                model_data = {"epoch": epoch, "optimizer_state_dict": deepcopy(self.optimizer.state_dict()), "loss": train_loss}
+                torch.save(model_data, self.model_dir + "/model_data_epoch" + str(epoch) + ".pt")
 
             torch.save(best_model_state_dict, self.model_dir + "/model.pt")
+            torch.save(best_model_data, self.model_dir + "/model_data.pt")
             torch.save(best_train_model_state_dict, self.model_dir + "/model_train.pt")
+            torch.save(best_train_model_data, self.model_dir + "/model_train_data.pt")
             elapsed_time = time.time() - start_time
             # Print elapsed time in minutes
             print(f"Elapsed time: {elapsed_time / 60:.2f} minutes")
diff --git a/utils/scripts/run.sh b/utils/scripts/run.sh
@@ -4,13 +4,13 @@
 SCRIPT_DIR="${CoverageControl_ws}/src/CoverageControl/python"
 
 # Set the parameters directory based on the environment variable
-PARAMS_DIR="${CoverageControl_ws}/lpac/params/"
+PARAMS_DIR="${CoverageControl_ws}/lpac_512/params/"
 
 # Define the parameter file names
 DATA_PARAMS_FILE="data_params.toml"
 # DATA_GEN_ALGORITHM="--algorithm CentralizedCVT"
 LEARNING_PARAMS_FILE="learning_params.toml"
-EVAL_PARAMS_FILE="eval.toml"
+EVAL_PARAMS_FILE="eval_multi.toml"
 
 # Function to print messages in red
 print_error() {
@@ -58,10 +58,10 @@ fi
 # Edit and execute process_data.sh
 
 # Running the data generation script
-run_command "python ${SCRIPT_DIR}/data_generation/data_generation.py ${PARAMS_DIR}/${DATA_PARAMS_FILE} ${DATA_GEN_ALGORITHM} --split True" "Data Generation"
+# run_command "python ${SCRIPT_DIR}/data_generation/data_generation.py ${PARAMS_DIR}/${DATA_PARAMS_FILE} ${DATA_GEN_ALGORITHM} --split True" "Data Generation"
 
 # Running the training script
-run_command "python ${SCRIPT_DIR}/training/train_lpac.py ${PARAMS_DIR}/${LEARNING_PARAMS_FILE} 1024" "Model Training"
+# run_command "python ${SCRIPT_DIR}/training/train_lpac.py ${PARAMS_DIR}/${LEARNING_PARAMS_FILE} 512" "Model Training"
 
 # Running the evaluation script
 run_command "python ${SCRIPT_DIR}/evaluators/eval.py ${PARAMS_DIR}/${EVAL_PARAMS_FILE}" "Model Evaluation"
