Merge pull request #16 from zhaiwenxi/dpa-adapt-resolve-pretrained-all-strategies

Dpa adapt resolve pretrained all strategies

Author: zhaiwenxi

doc/dpa_adapt/README.md

@@ -12,17 +12,6 @@ Installs `scikit-learn`, `dpdata`, `ase`, `rdkit`, and `e3nn` alongside DeePMD-k
 
 ## Quickstart
 
-Five lines to fine-tune and predict on CPU:
-
-```python
-from dpa_adapt import DPAFineTuner
-
-model = DPAFineTuner(pretrained="DPA-3.1-3M", strategy="frozen_sklearn", predictor="rf")
-model.fit(train_data="data/train", target_key="bandgap")
-preds = model.predict("data/test").predictions
-model.freeze("model.pth")
-```
-
 For a complete runnable example (QM9 HOMO–LUMO gap, ~5 min on CPU), see [`../../examples/dpa_adapt/`](../../examples/dpa_adapt/).
 
 ## Fine-tuning strategies
@@ -32,7 +21,7 @@ The strategy is the core choice. All four share the same pre-trained DPA backbon
 | Strategy         | Core Mechanism                                  | Target Data Size | Hardware     | Primary Use Case                          |
 | :--------------- | :---------------------------------------------- | :--------------- | :----------- | :---------------------------------------- |
 | `frozen_sklearn` | Frozen backbone + scikit-learn regressor        | Small (\<1k)     | CPU only     | Ultra-fast benchmarking & prototyping     |
-| `linear_probe`   | Frozen backbone + gradient-descent linear head  | Medium (1k–10k)  | CPU / GPU    | Balanced efficiency for linear properties |
+| `frozen_head`    | Frozen backbone + DeepMD property fitting head  | Medium (1k–10k)  | CPU / GPU    | Train only the property head while keeping the pretrained DPA backbone frozen |
 | `finetune`       | End-to-end full parameter fine-tuning           | Large (>10k)     | GPU required | Maximum accuracy on large datasets        |
 | `mft`            | Multi-task co-training (property + force field) | Small / low-data | GPU required | Mitigating representation collapse        |
 
@@ -46,9 +35,9 @@ model = DPAFineTuner(
 )
 model.fit(train_data="/data/train", target_key="homo")
 
-# linear_probe / finetune — same interface, different depth
+# frozen_head / finetune — same interface, different depth
 model = DPAFineTuner(
-    pretrained="DPA-3.1-3M", strategy="linear_probe", property_name="homo"
+    pretrained="DPA-3.1-3M", strategy="frozen_head", property_name="homo"
 )
 model.fit(train_data="/data/train", valid_data="/data/valid", target_key="homo")
 
@@ -111,7 +100,7 @@ model.predict(test_data, conditions={"temperature": T_test})
 # ConditionManager standardizes and concatenates values to the descriptor
 ```
 
-**linear_probe / finetune / mft** — place `fparam.npy` of shape `(nframes, fparam_dim)` in each `set.*/` directory alongside `coord.npy`, then declare the dimension at construction:
+**frozen_head / finetune / mft** — place `fparam.npy` of shape `(nframes, fparam_dim)` in each `set.*/` directory alongside `coord.npy`, then declare the dimension at construction:
 
 ```python
 model = DPAFineTuner(strategy="finetune", fparam_dim=2)
@@ -165,7 +154,7 @@ result = cross_validate(model, systems, label_key="energy", cv=5, group_by="form
 
 ```python
 from dpa_adapt import (
-    DPAFineTuner,  # fine-tune (strategies: frozen_sklearn, linear_probe, finetune, mft)
+    DPAFineTuner,  # fine-tune (strategies: frozen_sklearn, frozen_head, finetune, mft)
     DPAPredictor,  # inference from frozen bundles
     extract_descriptors,  # standalone descriptor extraction
     cross_validate,  # leak-proof cross-validation
@@ -196,7 +185,7 @@ X = extract_descriptors(
 
 | Command | Description |
 |---------|-------------|
-| `dpa-adapt fit` / `dpaad fit` | Fine-tune (`--strategy frozen_sklearn\|linear_probe\|finetune\|mft`) |
+| `dpa-adapt fit` / `dpaad fit` | Fine-tune (`--strategy frozen_sklearn\|frozen_head\|finetune\|mft`) |
 | `dpa-adapt predict` / `dpaad predict` | Predict with a frozen `.pth` bundle |
 | `dpa-adapt evaluate` / `dpaad evaluate` | Evaluate against stored labels |
 | `dpa-adapt extract-descriptors` / `dpaad extract-descriptors` | Extract pooled DPA descriptors to `.npy` |

dpa_adapt/cli.py

@@ -465,7 +465,7 @@ def get_parser() -> argparse.ArgumentParser:
     parser_fit.add_argument(
         "--strategy",
         default="frozen_sklearn",
-        choices=["frozen_sklearn", "linear_probe", "finetune", "mft"],
+        choices=["frozen_sklearn", "frozen_head", "finetune", "mft"],
     )
     parser_fit.add_argument(
         "--predictor", default="rf", choices=["rf", "linear", "ridge", "mlp"]
@@ -541,7 +541,7 @@ def get_parser() -> argparse.ArgumentParser:
         "--fparam-dim",
         type=int,
         default=0,
-        help="(linear_probe/finetune/mft) Dimensionality of per-frame condition "
+        help="(frozen_head/finetune/mft) Dimensionality of per-frame condition "
         "inputs (fparam). Requires set.*/fparam.npy in training data. Default: 0.",
     )
 

dpa_adapt/cv.py

@@ -305,7 +305,7 @@ def cross_validate(
     extracted **once** and a cheap sklearn head is trained per fold — even
     ``cv=5`` completes in seconds.
 
-    Training paradigms (``linear_probe`` / ``finetune`` / ``mft``)
+    Training paradigms (``frozen_head`` / ``finetune`` / ``mft``)
     are expensive: each fold re-trains a full DeepMD model.  To prevent
     accidental hour-long runs, *allow_expensive_cv* must be explicitly set
     to ``True`` for those strategies when *cv* is an integer >= 2.  Otherwise
@@ -529,7 +529,7 @@ def cross_validate(
             # Phase 2 will wire this to DPATrainer / MFTFineTuner.
             raise NotImplementedError(
                 "cross_validate for training paradigms "
-                "(linear_probe / finetune / mft) is not yet "
+                "(frozen_head / finetune / mft) is not yet "
                 "implemented. Use frozen_sklearn for now."
             )
 
@@ -578,7 +578,7 @@ def cross_validate(
 
 def _estimate_runtime(strategy: str, n_splits: int) -> str:
     per_run = {
-        "linear_probe": "~5-15 min/run",
+        "frozen_head": "~5-15 min/run",
         "finetune": "~10-30 min/run",
         "mft": "~20-60 min/run",
     }.get(strategy, "unknown")

dpa_adapt/data/type_map.py

@@ -34,8 +34,10 @@ def read_checkpoint_type_map(
     """
     from dpa_adapt._backend import (
         load_torch_file,
+        resolve_pretrained_path,
     )
 
+    pretrained = resolve_pretrained_path(pretrained)
     sd = load_torch_file(pretrained)
     if "model" in sd:
         sd = sd["model"]

dpa_adapt/finetuner.py

@@ -510,7 +510,7 @@ class DPAFineTuner:
                           descriptors once, pool, and fit a scikit-learn
                           regressor (Ridge, KRR, or MLP).  No GPU needed;
                           fastest for small datasets.
-    ``linear_probe``      Freeze the backbone, train only a neural property
+    ``frozen_head``      Freeze the backbone, train only a neural property
                           fitting net via ``dp --pt train``.
     ``finetune``          Fine-tune the full network (descriptor + fitting
                           net) end-to-end via ``dp --pt train``.
@@ -536,7 +536,7 @@ class DPAFineTuner:
     seed : int
         Random seed for the head or for full training.
     strategy : str
-        ``"frozen_sklearn"`` (default), ``"linear_probe"``, ``"finetune"``,
+        ``"frozen_sklearn"`` (default), ``"frozen_head"``, ``"finetune"``,
         or ``"mft"``.
 
     property_name : str
@@ -560,7 +560,7 @@ class DPAFineTuner:
     loss_function : str
         ``"mse"`` or ``"smooth_mae"`` (training paradigms).
     fparam_dim : int
-        (linear_probe / finetune / mft only) Dimensionality of per-frame
+        (frozen_head / finetune / mft only) Dimensionality of per-frame
         condition inputs (e.g. temperature, pressure). Requires
         set.*/fparam.npy of shape (n_frames, fparam_dim) in every
         training system. Default 0 (disabled).
@@ -590,7 +590,7 @@ class DPAFineTuner:
     _VALID_POOLING = {"mean", "sum", "mean+std", "mean+std+max+min"}
     _VALID_STRATEGIES = {
         "frozen_sklearn",
-        "linear_probe",
+        "frozen_head",
         "finetune",
         "mft",
     }
@@ -807,7 +807,7 @@ def _resolve_type_maps(self, train_data) -> list[str]:
         return tm
 
     # -------------------------------------------------------------------
-    # Training-paradigm fit (linear_probe / finetune)
+    # Training-paradigm fit (frozen_head / finetune)
     # -------------------------------------------------------------------
 
     def _fit_training(self, train_data, valid_data, type_map):
@@ -816,7 +816,7 @@ def _fit_training(self, train_data, valid_data, type_map):
             DPATrainer,
         )
 
-        freeze = self.strategy == "linear_probe"
+        freeze = self.strategy == "frozen_head"
         trainer = DPATrainer(
             pretrained=self.pretrained,
             init_branch=self.init_branch,
@@ -860,7 +860,7 @@ def fit(
         """Train the model.
 
         *frozen_sklearn* (default): extract descriptors, fit sklearn head.
-        *linear_probe* / *finetune*: run ``dp --pt train``.
+        *frozen_head* / *finetune*: run ``dp --pt train``.
         *mft*: multi-task fine-tuning (property head + force-field head).
 
         Parameters

dpa_adapt/mft.py

@@ -5,6 +5,11 @@
 import subprocess
 import sys
 
+from dpa_adapt._backend import (
+    load_torch_file,
+    resolve_pretrained_path,
+)
+
 
 class MFTFineTuner:
     """
@@ -129,7 +134,7 @@ def __init__(
                 f"fparam_dim must be a non-negative int; got {fparam_dim!r}."
             )
 
-        self.pretrained = pretrained
+        self.pretrained = resolve_pretrained_path(pretrained)
         self.aux_branch = aux_branch
         self.aux_prob = aux_prob
         self.aux_type_map = aux_type_map
@@ -187,9 +192,7 @@ def _read_fitting_net_from_ckpt(pretrained, aux_branch):
         checkpoint. Raises ValueError listing available branches if
         ``aux_branch`` isn't present.
         """
-        import torch
-
-        sd = torch.load(pretrained, map_location="cpu", weights_only=False)
+        sd = load_torch_file(resolve_pretrained_path(pretrained))
         try:
             model_dict = sd["model"]["_extra_state"]["model_params"]["model_dict"]
         except (KeyError, TypeError) as e:

dpa_adapt/trainer.py

@@ -29,6 +29,10 @@
 import re
 import subprocess
 
+from dpa_adapt._backend import (
+    resolve_pretrained_path,
+)
+
 _LOG = logging.getLogger("dpa_adapt.trainer")
 
 
@@ -190,8 +194,10 @@ def __init__(
                 "LP requires a pretrained checkpoint to freeze. "
                 "Set freeze_backbone=False for Scratch, or pass a pretrained ckpt."
             )
-        if pretrained is not None and not os.path.isfile(pretrained):
-            raise ValueError(f"pretrained checkpoint not found: {pretrained!r}.")
+        if pretrained is not None:
+            pretrained = resolve_pretrained_path(pretrained)
+            if not os.path.isfile(pretrained):
+                raise ValueError(f"pretrained checkpoint not found: {pretrained!r}.")
         if not isinstance(property_name, str) or not property_name.isidentifier():
             raise ValueError(
                 f"property_name must be a valid Python identifier "

pyproject.toml

@@ -98,6 +98,7 @@ dpa-adapt = [
     "scikit-learn",
     "dpdata",
     "torch",
+    "ase",
     "rdkit",
     "e3nn",
 ]

source/tests/dpa_adapt/test_finetuner_strategies.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Tests for DPAFineTuner training-paradigm strategies
-(linear_probe / finetune).
+(frozen_head / finetune).
 
 Mock ``dp --pt train`` via ``subprocess.run``; verify:
 - Correct DPATrainer params per strategy
@@ -188,7 +188,7 @@ def test_resolve_type_maps_from_checkpoint(self, monkeypatch, tmp_path):
         systems = _make_system_dirs(tmp_path)
         m = DPAFineTuner(
             pretrained="/fake.pt",
-            strategy="linear_probe",
+            strategy="frozen_head",
             init_branch="SPICE2",
         )
         tm = m._resolve_type_maps(systems)
@@ -241,7 +241,7 @@ def _mock_torch(self, monkeypatch, tmp_path):
     @pytest.mark.parametrize(
         "strategy,expect_freeze,expect_tm_len",
         [
-            ("linear_probe", True, 8),
+            ("frozen_head", True, 8),
             ("finetune", False, 8),
         ],
     )
@@ -287,7 +287,7 @@ def test_config_type_map_nonempty(
         )
         assert tm != [], "type_map is empty — would cause CUDA gather out-of-bounds"
 
-    @pytest.mark.parametrize("strategy", ["linear_probe", "finetune"])
+    @pytest.mark.parametrize("strategy", ["frozen_head", "finetune"])
     def test_strategy_to_trainer_params(self, tmp_path, strategy):
         """Each strategy produces correct DPATrainer freeze_backbone / pretrained."""
         out_dir = tmp_path / "out"
@@ -317,7 +317,7 @@ def test_strategy_to_trainer_params(self, tmp_path, strategy):
         assert fn["intensive"] is True
 
         # LP must freeze backbone
-        if strategy == "linear_probe":
+        if strategy == "frozen_head":
             assert cfg["model"]["descriptor"]["trainable"] is False
         else:
             assert cfg["model"]["descriptor"]["trainable"] is True