implement training on per atom energies by modifying loss modules

src/metatensor/models/cli/conf/architecture/experimental.alchemical_model.yaml

@@ -18,3 +18,6 @@ training:
   learning_rate: 0.001
   log_interval: 10
   checkpoint_interval: 25
+  per_atom_targets: []  # this specifies whether the model should be trained on a per-atom loss.
+                        # In that case, the logger will also output per-atom metrics for that
+                        # target. In any case, the final summary will be per-structure.

src/metatensor/models/cli/conf/architecture/experimental.soap_bpnn.yaml

@@ -25,3 +25,6 @@ training:
   learning_rate: 0.001
   log_interval: 10
   checkpoint_interval: 25
+  per_atom_targets: []  # this specifies whether the model should be trained on a per-atom loss.
+                        # In that case, the logger will also output per-atom metrics for that
+                        # target. In any case, the final summary will be per-structure.

src/metatensor/models/cli/eval.py

@@ -103,8 +103,10 @@ def _eval_targets(model, dataset: Union[_BaseDataset, torch.utils.data.Subset])
     # Compute the RMSEs:
     aggregated_info: Dict[str, Tuple[float, int]] = {}
     for batch in dataloader:
-        systems, targets = batch
-        _, info = compute_model_loss(loss_fn, model, systems, targets)
+
+        structures, targets = batch
+        _, info = compute_model_loss(loss_fn, model, structures, targets, [])
+
         aggregated_info = update_aggregated_info(aggregated_info, info)
     finalized_info = finalize_aggregated_info(aggregated_info)
 

src/metatensor/models/experimental/alchemical_model/train.py

@@ -207,9 +207,13 @@ def train(
         train_loss = 0.0
         for batch in train_dataloader:
             optimizer.zero_grad()
+
             systems, targets = batch
             assert len(systems[0].known_neighbors_lists()) > 0
-            loss, info = compute_model_loss(loss_fn, model, systems, targets)
+            loss, info = compute_model_loss(
+                loss_fn, model, systems, targets, hypers_training["per_atom_targets"]
+            )
+
             train_loss += loss.item()
             loss.backward()
             optimizer.step()
@@ -220,7 +224,11 @@ def train(
         for batch in validation_dataloader:
             systems, targets = batch
             # TODO: specify that the model is not training here to save some autograd
-            loss, info = compute_model_loss(loss_fn, model, systems, targets)
+
+            loss, info = compute_model_loss(
+                loss_fn, model, systems, targets, hypers_training["per_atom_targets"]
+            )
+
             validation_loss += loss.item()
             aggregated_validation_info = update_aggregated_info(
                 aggregated_validation_info, info

src/metatensor/models/experimental/soap_bpnn/train.py

@@ -182,8 +182,12 @@ def train(
         train_loss = 0.0
         for batch in train_dataloader:
             optimizer.zero_grad()
+
             systems, targets = batch
-            loss, info = compute_model_loss(loss_fn, model, systems, targets)
+            loss, info = compute_model_loss(
+                loss_fn, model, systems, targets, hypers_training["per_atom_targets"]
+            )
+
             train_loss += loss.item()
             loss.backward()
             optimizer.step()
@@ -194,7 +198,11 @@ def train(
         for batch in validation_dataloader:
             systems, targets = batch
             # TODO: specify that the model is not training here to save some autograd
-            loss, info = compute_model_loss(loss_fn, model, systems, targets)
+
+            loss, info = compute_model_loss(
+                loss_fn, model, systems, targets, hypers_training["per_atom_targets"]
+            )
+
             validation_loss += loss.item()
             aggregated_validation_info = update_aggregated_info(
                 aggregated_validation_info, info

src/metatensor/models/utils/compute_loss.py

@@ -29,16 +29,22 @@ def compute_model_loss(
     model: Union[torch.nn.Module, torch.jit._script.RecursiveScriptModule],
     systems: List[System],
     targets: Dict[str, TensorMap],
+    per_atom_targets: List[str],
 ) -> Tuple[torch.Tensor, Dict[str, Tuple[float, int]]]:
     """
-    Compute the loss of a model on a set of targets.
+    Compute the loss of a model on a set of targets, with an option to treat
+    specifed targets on a per atom basis. This implies that when some such
+    targets are specified, their contribution to the loss will accordingly be on
+    a per atom basis.
 
     :param loss: The loss function to use.
     :param model: The model to use. This can either be a model in training
         (``torch.nn.Module``) or an exported model
         (``torch.jit._script.RecursiveScriptModule``).
     :param systems: The systems to use.
     :param targets: The targets to use.
+    :param per_atom_targets: The targets that should be treated on a per atom
+        basis during loss calculation.
 
     :returns: The loss as a scalar `torch.Tensor`.
     """
@@ -174,8 +180,37 @@ def compute_model_loss(
         else:
             pass
 
+    # Averaging by number of atoms for per atom targets
+    num_atoms = torch.tensor([len(s) for s in systems], device=device).unsqueeze(-1)
+
+    new_model_outputs = model_outputs.copy()
+    new_targets = targets.copy()
+
+    for pa_target in per_atom_targets:
+
+        # Update predictions
+        cur_model_block = new_model_outputs[pa_target].block()
+        new_model_block = _average_by_num_atoms(cur_model_block, num_atoms)
+
+        # Update targets
+        cur_target_block = new_targets[pa_target].block()
+        new_target_block = _average_by_num_atoms(cur_target_block, num_atoms)
+
+        new_model_tensor = TensorMap(
+            keys=new_model_outputs[pa_target].keys,
+            blocks=[new_model_block],
+        )
+
+        new_target_tensor = TensorMap(
+            keys=new_targets[pa_target].keys,
+            blocks=[new_target_block],
+        )
+
+        new_model_outputs[pa_target] = new_model_tensor
+        new_targets[pa_target] = new_target_tensor
+
     # Compute and return the loss and associated info:
-    return loss(model_outputs, targets)
+    return loss(new_model_outputs, new_targets)
 
 
 def _position_gradients_to_block(gradients_list):
@@ -275,3 +310,19 @@ def _get_model_outputs(
         return model(
             systems, {key: _get_capabilities(model).outputs[key] for key in targets}
         )
+
+
+def _average_by_num_atoms(block: TensorBlock, num_atoms: torch.Tensor) -> TensorBlock:
+    """Taking the average values per atom of a `TensorBlock`."""
+
+    new_values = block.values / num_atoms
+    new_block = TensorBlock(
+        values=new_values,
+        samples=block.samples,
+        components=block.components,
+        properties=block.properties,
+    )
+    for param, gradient in block.gradients():
+        new_block.add_gradient(param, gradient)
+
+    return new_block

tests/utils/test_compute_loss.py

@@ -94,9 +94,22 @@ def test_compute_model_loss():
         ),
     }
 
-    compute_model_loss(
+    loss, info = compute_model_loss(
         loss_fn,
         model,
         systems,
         targets,
+        [],
     )
+
+    per_atom_targets = ["energy"]
+
+    per_atom_loss, info = compute_model_loss(
+        loss_fn,
+        model,
+        systems,
+        targets,
+        per_atom_targets,
+    )
+
+    assert loss > per_atom_loss