adding lr settings to allowed settings + enabling checkpointing

tulga-rdn · Jan 30, 2025 · a395f81 · a395f81
1 parent e62386c
commit a395f81
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/default_hypers/default_hypers.yaml b/default_hypers/default_hypers.yaml
@@ -59,7 +59,7 @@ FITTING_SCHEME:
 
 LONG_RANGE_SETTINGS: # torchPME settings
   FULL_NEIGHBOUR_LIST: False
-  PREFACTOR: 14.
+  PREFACTOR: 1.0
   CHARGE_CHANNELS: 1
   ATOMIC_SMEARING: 1.
   LR_WAVELENGTH: 0.5

diff --git a/src/estimate_error.py b/src/estimate_error.py
@@ -91,6 +91,8 @@ def main():
     FITTING_SCHEME = hypers.FITTING_SCHEME
 
     ARCHITECTURAL_HYPERS = hypers.ARCHITECTURAL_HYPERS
+
+    LONG_RANGE_SETTINGS = hypers.LONG_RANGE_SETTINGS
 
     # set_reproducibility(FITTING_SCHEME.RANDOM_SEED, FITTING_SCHEME.CUDA_DETERMINISTIC)
 
@@ -122,7 +124,8 @@ def main():
 
     if hypers.UTILITY_FLAGS.CALCULATION_TYPE == "mlip":
         model = PETMLIPWrapper(
-            model, hypers.MLIP_SETTINGS.USE_ENERGIES, hypers.MLIP_SETTINGS.USE_FORCES
+            model, hypers.MLIP_SETTINGS.USE_ENERGIES, hypers.MLIP_SETTINGS.USE_FORCES,
+            LONG_RANGE_SETTINGS, device
         )
 
     if FITTING_SCHEME.MULTI_GPU and torch.cuda.is_available():

diff --git a/src/train_model.py b/src/train_model.py
@@ -125,6 +125,7 @@ def fit_pet(
     save_hypers(hypers, f"{output_dir}/{NAME_OF_CALCULATION}/hypers_used.yaml")
 
     print("Convering structures to PyG graphs...", flush=True)
+    print(len(train_structures))
 
     train_graphs = get_pyg_graphs(
         train_structures,
@@ -557,6 +558,9 @@ def main():
     parser.add_argument(
         "name_of_calculation", help="Name of this calculation", type=str
     )
+    parser.add_argument(
+        "--checkpoint", help="checkpoint path", type=str, default=None
+    )
     parser.add_argument("--gpu_id", help="ID of the GPU to use", type=int, default=0)
     args = parser.parse_args()
 
@@ -572,7 +576,7 @@ def main():
 
     name_of_calculation = args.name_of_calculation
 
-    output_dir = "results"
+    output_dir = "/scratch/izar/sodjarga/results"
 
     hypers_dict = hypers_to_dict(hypers)
     fit_pet(
@@ -582,6 +586,7 @@ def main():
         name_of_calculation,
         device,
         output_dir,
+        args.checkpoint
     )