[fix] [torch_training] Resuming from checkpoints now first evaluates the entire data

The previous commit introduced sampling policies that can intentionally result in not sampling all of the data or sampling some data points multiple times. This can lead to checkpoint files that do not jhave reproducible error metrics. Hence, after resuming from a checkpoint, the error metrics need to reevaluated according to the used sampling policy.

Author: alexurba

src/aenet/torch_training/trainer.py

@@ -1945,6 +1945,79 @@ def train(
             == "error_weighted"
         )
 
+        def _replace_latest_metrics_with_full_pass() -> None:
+            history = self._metrics.get_history()
+            if len(history.get("train_energy_rmse", [])) == 0:
+                return
+
+            final_train_loader = DataLoader(
+                train_ds,
+                batch_size=batch_size,
+                shuffle=False,
+                collate_fn=_collate_fn,
+                **eval_dl_kwargs,
+            )
+            (
+                final_train_energy_rmse,
+                final_train_energy_mae,
+                final_train_force_rmse,
+                _,
+                _,
+            ) = training_loop.run_epoch(
+                loader=final_train_loader,
+                optimizer=None,
+                alpha=alpha,
+                atomic_energies_by_index=atomic_energies_by_index,
+                train=False,
+                show_batch_progress=False,
+                force_scale_unbiased=bool(
+                    getattr(config, "force_scale_unbiased", False)
+                ),
+                collect_structure_scores=False,
+            )
+
+            final_test_energy_rmse = float("nan")
+            final_test_energy_mae = float("nan")
+            final_test_force_rmse = float("nan")
+            if test_ds is not None:
+                final_test_loader = DataLoader(
+                    test_ds,
+                    batch_size=batch_size,
+                    shuffle=False,
+                    collate_fn=_collate_fn,
+                    **eval_dl_kwargs,
+                )
+                (
+                    final_test_energy_rmse,
+                    final_test_energy_mae,
+                    final_test_force_rmse,
+                    _,
+                    _,
+                ) = training_loop.run_epoch(
+                    loader=final_test_loader,
+                    optimizer=None,
+                    alpha=alpha,
+                    atomic_energies_by_index=atomic_energies_by_index,
+                    train=False,
+                    show_batch_progress=False,
+                    force_scale_unbiased=bool(
+                        getattr(config, "force_scale_unbiased", False)
+                    ),
+                    collect_structure_scores=False,
+                )
+
+            self._metrics.replace_latest(
+                train_energy_rmse=float(final_train_energy_rmse),
+                train_energy_mae=float(final_train_energy_mae),
+                train_force_rmse=float(final_train_force_rmse),
+                test_energy_rmse=float(final_test_energy_rmse),
+                test_energy_mae=float(final_test_energy_mae),
+                test_force_rmse=float(final_test_force_rmse),
+            )
+
+        if start_epoch > 0:
+            _replace_latest_metrics_with_full_pass()
+
         for epoch in range(start_epoch, end_epoch):
             t0 = time.time()
             save_best_checkpoint = False
@@ -2088,70 +2161,7 @@ def train(
             pbar.close()
 
         if n_epochs > 0:
-            final_train_loader = DataLoader(
-                train_ds,
-                batch_size=batch_size,
-                shuffle=False,
-                collate_fn=_collate_fn,
-                **eval_dl_kwargs,
-            )
-            (
-                final_train_energy_rmse,
-                final_train_energy_mae,
-                final_train_force_rmse,
-                _,
-                _,
-            ) = training_loop.run_epoch(
-                loader=final_train_loader,
-                optimizer=None,
-                alpha=alpha,
-                atomic_energies_by_index=atomic_energies_by_index,
-                train=False,
-                show_batch_progress=False,
-                force_scale_unbiased=bool(
-                    getattr(config, "force_scale_unbiased", False)
-                ),
-                collect_structure_scores=False,
-            )
-
-            final_test_energy_rmse = float("nan")
-            final_test_energy_mae = float("nan")
-            final_test_force_rmse = float("nan")
-            if test_ds is not None:
-                final_test_loader = DataLoader(
-                    test_ds,
-                    batch_size=batch_size,
-                    shuffle=False,
-                    collate_fn=_collate_fn,
-                    **eval_dl_kwargs,
-                )
-                (
-                    final_test_energy_rmse,
-                    final_test_energy_mae,
-                    final_test_force_rmse,
-                    _,
-                    _,
-                ) = training_loop.run_epoch(
-                    loader=final_test_loader,
-                    optimizer=None,
-                    alpha=alpha,
-                    atomic_energies_by_index=atomic_energies_by_index,
-                    train=False,
-                    show_batch_progress=False,
-                    force_scale_unbiased=bool(
-                        getattr(config, "force_scale_unbiased", False)
-                    ),
-                    collect_structure_scores=False,
-                )
-
-            self._metrics.replace_latest(
-                train_energy_rmse=float(final_train_energy_rmse),
-                train_energy_mae=float(final_train_energy_mae),
-                train_force_rmse=float(final_train_force_rmse),
-                test_energy_rmse=float(final_test_energy_rmse),
-                test_energy_mae=float(final_test_energy_mae),
-                test_force_rmse=float(final_test_force_rmse),
-            )
+            _replace_latest_metrics_with_full_pass()
 
         # Store optimizer and config for later saving
         self._optimizer = optimizer