From a59ff7d66b06543b1c3533e9d37038b07dc0e191 Mon Sep 17 00:00:00 2001
From: frostedoyster <bigi.f@libero.it>
Date: Fri, 4 Oct 2024 18:42:21 +0200
Subject: [PATCH] Add evaluation timings

---
 src/metatrain/cli/eval.py    | 24 ++++++++++++++++++++++++
 tests/cli/test_eval_model.py |  2 ++
 2 files changed, 26 insertions(+)

diff --git a/src/metatrain/cli/eval.py b/src/metatrain/cli/eval.py
index 93adb162..6b308cae 100644
--- a/src/metatrain/cli/eval.py
+++ b/src/metatrain/cli/eval.py
@@ -1,6 +1,7 @@
 import argparse
 import itertools
 import logging
+import time
 from pathlib import Path
 from typing import Dict, List, Optional, Union
 
@@ -194,6 +195,10 @@ def _eval_targets(
     if return_predictions:
         all_predictions = []
 
+    # Set up timings:
+    total_time = 0.0
+    timings_per_atom = []
+
     # Evaluate the model
     for batch in dataloader:
         systems, batch_targets = batch
@@ -202,6 +207,9 @@ def _eval_targets(
             key: value.to(dtype=dtype, device=device)
             for key, value in batch_targets.items()
         }
+
+        start_time = time.time()
+
         batch_predictions = evaluate_model(
             model,
             systems,
@@ -209,6 +217,11 @@ def _eval_targets(
             is_training=False,
             check_consistency=check_consistency,
         )
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        end_time = time.time()
+
         batch_predictions = average_by_num_atoms(
             batch_predictions, systems, per_structure_keys=[]
         )
@@ -219,6 +232,10 @@ def _eval_targets(
         if return_predictions:
             all_predictions.append(batch_predictions)
 
+        time_taken = end_time - start_time
+        total_time += time_taken
+        timings_per_atom.append(time_taken / sum(len(system) for system in systems))
+
     # Finalize the RMSEs
     rmse_values = rmse_accumulator.finalize(not_per_atom=["positions_gradients"])
     # print the RMSEs with MetricLogger
@@ -229,6 +246,13 @@ def _eval_targets(
     )
     metric_logger.log(rmse_values)
 
+    # Log timings
+    mean_time_per_atom = sum(timings_per_atom) / len(timings_per_atom)
+    logger.info(
+        f"evaluation time: {total_time:.2f} s "
+        f"[{1000.0*mean_time_per_atom:.2f} ms per atom]"
+    )
+
     if return_predictions:
         # concatenate the TensorMaps
         all_predictions_joined = _concatenate_tensormaps(all_predictions)
diff --git a/tests/cli/test_eval_model.py b/tests/cli/test_eval_model.py
index 7aaa2c2b..7eda92af 100644
--- a/tests/cli/test_eval_model.py
+++ b/tests/cli/test_eval_model.py
@@ -68,6 +68,8 @@ def test_eval(monkeypatch, tmp_path, caplog, model_name, options):
     log = "".join([rec.message for rec in caplog.records])
     assert "energy RMSE (per atom)" in log
     assert "dataset with index" not in log
+    assert "evaluation time" in log
+    assert "ms per atom" in log
 
     # Test file is written predictions
     frames = ase.io.read("foo.xyz", ":")