added option to train AtomRef in trainer

chgnet/trainer/trainer.py

@@ -200,17 +200,24 @@ def train(
         test_loader: DataLoader | None = None,
         save_dir: str | None = None,
         save_test_result: bool = False,
+        train_composition_model: bool = False,
     ) -> None:
         """Train the model using torch data_loaders.
 
         Args:
             train_loader (DataLoader): train loader to update CHGNet weights
             val_loader (DataLoader): val loader to test accuracy after each epoch
             test_loader (DataLoader):  test loader to test accuracy at end of training.
-                Can be None. Default = None.
+                Can be None.
+                Default = None
             save_dir (str): the dir name to save the trained weights
                 Default = None
             save_test_result (bool): whether to save the test set prediction in a json file
+            train_composition_model (bool): whether to train the composition model
+                (AtomRef), this is suggested when the fine-tuning dataset has large
+                elemental energy shift from the pretrained CHGNet, which typically comes
+                from different DFT pseudo-potentials.
+                Default = False
         """
         if self.model is None:
             raise ValueError("Model needs to be initialized")
@@ -223,6 +230,10 @@ def train(
         print(f"training targets: {self.targets}")
         self.model.to(self.device)
 
+        # Turn composition model training on / off
+        for param in self.model.composition_model.parameters():
+            param.requires_grad = train_composition_model
+
         for epoch in range(self.starting_epoch, self.epochs):
             # train
             train_mae = self._train(train_loader, epoch)

tests/test_trainer.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import numpy as np
+import torch
 from pymatgen.core import Lattice, Structure
 
 from chgnet.data.dataset import StructureData, get_train_val_test_loader
@@ -52,3 +53,42 @@ def test_trainer(tmp_path) -> None:
     for prefix in ("epoch", "bestE", "bestF"):
         n_matches = sum(file.name.startswith(prefix) for file in output_files)
         assert n_matches == 1
+
+
+def test_trainer_composition_model(tmp_path) -> None:
+    chgnet = CHGNet.load()
+    for param in chgnet.composition_model.parameters():
+        assert param.requires_grad is False
+    train_loader, val_loader, test_loader = get_train_val_test_loader(
+        data, batch_size=16, train_ratio=0.9, val_ratio=0.05
+    )
+    trainer = Trainer(
+        model=chgnet,
+        targets="efsm",
+        optimizer="Adam",
+        criterion="MSE",
+        learning_rate=1e-2,
+        epochs=5,
+    )
+    dir_name = "test_tmp_dir2"
+    test_dir = tmp_path / dir_name
+    initial_weights = chgnet.composition_model.state_dict()["fc.weight"].clone()
+    trainer.train(
+        train_loader, val_loader, save_dir=test_dir, train_composition_model=True
+    )
+    for param in chgnet.composition_model.parameters():
+        assert param.requires_grad is True
+
+    output_files = list(test_dir.iterdir())
+    weights_path = next(file for file in output_files if file.name.startswith("epoch"))
+    new_chgnet = CHGNet.from_file(weights_path)
+    for param in new_chgnet.composition_model.parameters():
+        assert param.requires_grad is False
+    comparison = (
+        new_chgnet.composition_model.state_dict()["fc.weight"] == initial_weights
+    )
+    expect = torch.ones_like(comparison)
+    # Only Na and Cl should have updated
+    expect[0][10] = 0
+    expect[0][16] = 0
+    assert torch.all(comparison == expect)