Merge pull request #9 from zirenjin/master

dpa_tools quickstart demo bundle, convert() glob multi-match, docs cleanup

Author: zhaiwenxi

README.md

@@ -107,7 +107,7 @@ from deepmd.dpa_tools import DPAFineTuner
 
 model = DPAFineTuner(pretrained="DPA-3.1-3M", strategy="frozen_sklearn", predictor="rf")
 model.fit(train_data="data/train", target_key="bandgap")  # fine-tune on your labeled structures
-model.predict("data/new_structures")                      # predict for new structures
+preds = model.predict("data/new_structures").predictions   # predict for new structures
 ```
 
 The same workflow is also available from the command line as `dp dpa fit` / `dp dpa predict`.

deepmd/dpa_tools/demo/fit_evaluate.py

@@ -1,14 +1,17 @@
 #!/usr/bin/env python3
-"""Fit a frozen DPA descriptor + Ridge regressor on the quickstart demo data.
+"""Fit a pretrained DPA descriptor on the quickstart QM9 demo data.
 
-Requires the DPA-3.1-3M pretrained checkpoint.  Provide it via ``--model`` or
-set the ``DPA_MODEL_PATH`` environment variable.
+All four fine-tuning strategies are demonstrated.  The default
+(``frozen_sklearn``) runs on CPU in under 5 minutes and requires no GPU.
+``linear_probe``, ``finetune``, and ``mft`` use ``dp --pt train`` under the
+hood and need a GPU to finish in reasonable time.
 
 Usage (from the demo directory)::
 
     python fit_evaluate.py --model /path/to/DPA-3.1-3M.pt
+    python fit_evaluate.py --model /path/to/DPA-3.1-3M.pt --strategy finetune
 
-(or set the ``DPA_MODEL_PATH`` environment variable instead of ``--model``).
+Set ``DPA_MODEL_PATH`` instead of ``--model`` to avoid typing it every time.
 """
 
 from __future__ import annotations
@@ -29,90 +32,230 @@
 FROZEN_MODEL_PATH = HERE / "frozen_model.pth"
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        prog="dp dpa fit",
-        description="Quickstart: fit frozen DPA descriptor + Ridge on QM9 HOMO-LUMO gap.",
-    )
-    parser.add_argument(
-        "--model",
-        default=None,
-        help="Path to DPA-3.1-3M.pt checkpoint.  Falls back to $DPA_MODEL_PATH.",
-    )
-    args = parser.parse_args()
+# ---------------------------------------------------------------------------
+# helpers
+# ---------------------------------------------------------------------------
+
 
-    # --- resolve model path ---
-    model_path = args.model or os.environ.get("DPA_MODEL_PATH")
-    if not model_path:
+def _resolve_model(args: argparse.Namespace) -> str:
+    path = args.model or os.environ.get("DPA_MODEL_PATH")
+    if not path:
         print(
-            "error: DPA-3.1-3M checkpoint not specified.\n"
-            "  Provide it via --model or set the DPA_MODEL_PATH environment variable.\n"
+            "error: DPA checkpoint not specified.\n"
+            "  Provide it via --model or set $DPA_MODEL_PATH.\n"
             "  Example: python fit_evaluate.py --model /path/to/DPA-3.1-3M.pt",
             file=sys.stderr,
         )
         sys.exit(1)
-
-    if not Path(model_path).is_file():
-        print(f"error: model file not found: {model_path}", file=sys.stderr)
+    if not Path(path).is_file():
+        print(f"error: model file not found: {path}", file=sys.stderr)
         sys.exit(1)
+    return path
 
-    print(f"Model checkpoint: {model_path}")
 
-    # --- verify data ---
-    if not TRAIN_DIR.is_dir():
-        print(
-            f"error: training data not found at {TRAIN_DIR}\n"
-            "  Run scripts/prepare_data.py first.",
-            file=sys.stderr,
-        )
-        sys.exit(1)
-    if not TEST_DIR.is_dir():
-        print(
-            f"error: test data not found at {TEST_DIR}\n"
-            "  Run scripts/prepare_data.py first.",
-            file=sys.stderr,
-        )
-        sys.exit(1)
+def _verify_data() -> None:
+    for name, d in [("train", TRAIN_DIR), ("test", TEST_DIR)]:
+        if not d.is_dir():
+            print(
+                f"error: {name} data not found at {d}\n"
+                "  Run scripts/prepare_data.py first.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+
+
+def _load_labels() -> tuple[np.ndarray, np.ndarray]:
+    train = np.load(str(TRAIN_LABELS_PATH)).astype(np.float32)
+    test = np.load(str(TEST_LABELS_PATH)).astype(np.float32)
+    return train, test
+
+
+def _print_metrics(metrics, label: str = "") -> None:
+    tag = f" [{label}]" if label else ""
+    print()
+    print("=" * 50)
+    print(f"MAE{tag}  : {metrics.mae:.4f} eV")
+    print(f"R²{tag}   : {metrics.r2:.4f}")
+    print(f"RMSE{tag} : {metrics.rmse:.4f} eV")
+    print(f"N{tag}    : {metrics.predictions.shape[0]}")
+    print("=" * 50)
 
-    # --- load labels ---
-    train_labels = np.load(str(TRAIN_LABELS_PATH)).astype(np.float32)
-    test_labels = np.load(str(TEST_LABELS_PATH)).astype(np.float32)
 
-    # --- build model ---
+# ---------------------------------------------------------------------------
+# Strategy 1 — frozen_sklearn (default, CPU)
+# ---------------------------------------------------------------------------
+
+
+def demo_frozen_sklearn(model_path: str, train_labels: np.ndarray) -> None:
+    """Freeze the DPA backbone, extract descriptors once, fit a sklearn Ridge.
+
+    Fastest iteration.  No GPU, no ``dp train`` subprocess.  The frozen
+    bundle (``.pth``) is portable and can be loaded with ``DPAPredictor``.
+    """
     from deepmd.dpa_tools import DPAFineTuner
 
     model = DPAFineTuner(
         pretrained=model_path,
         model_branch="Domains_Drug",
-        pooling="mean",
-        predictor="linear",
+        strategy="frozen_sklearn",
+        predictor="linear",   # "linear" (Ridge) | "rf" | "mlp"
+        pooling="mean",        # "mean" | "sum" | "mean+std" | "mean+std+max+min"
         seed=42,
     )
 
-    # --- fit ---
-    print("Fitting …")
+    print("frozen_sklearn — fitting …")
     model.fit(
         train_data=str(TRAIN_DIR),
         labels=train_labels,
         target_key="gap",
     )
 
-    # --- evaluate ---
-    print("Evaluating …")
+    print("frozen_sklearn — evaluating …")
     metrics = model.evaluate(data=str(TEST_DIR))
+    _print_metrics(metrics, "frozen_sklearn")
 
-    print()
-    print("=" * 50)
-    print(f"MAE  : {metrics.mae:.4f} eV")
-    print(f"R²   : {metrics.r2:.4f}")
-    print(f"RMSE : {metrics.rmse:.4f} eV")
-    print(f"N    : {metrics.predictions.shape[0]}")
-    print("=" * 50)
-
-    # --- freeze ---
     out = model.freeze(str(FROZEN_MODEL_PATH))
     print(f"Frozen model → {out}")
 
 
+# ---------------------------------------------------------------------------
+# Strategy 2 — linear_probe (GPU recommended)
+# ---------------------------------------------------------------------------
+
+
+def demo_linear_probe(model_path: str) -> None:
+    """Freeze the DPA backbone, train only a neural property fitting net.
+
+    Uses ``dp --pt train --finetune`` under the hood.  A GPU is recommended.
+    """
+    from deepmd.dpa_tools import DPAFineTuner
+
+    model = DPAFineTuner(
+        pretrained=model_path,
+        strategy="linear_probe",
+        property_name="gap",
+        task_dim=1,
+        intensive=True,
+        output_dir=str(HERE / "output_lp"),
+    )
+
+    print("linear_probe — fitting (dp --pt train) …")
+    model.fit(
+        train_data=str(TRAIN_DIR),
+        valid_data=str(TEST_DIR),
+        target_key="gap",
+    )
+
+    print("linear_probe — evaluating …")
+    metrics = model.evaluate(data=str(TEST_DIR))
+    _print_metrics(metrics, "linear_probe")
+
+
+# ---------------------------------------------------------------------------
+# Strategy 3 — finetune (GPU recommended)
+# ---------------------------------------------------------------------------
+
+
+def demo_finetune(model_path: str) -> None:
+    """Load the pretrained backbone and fine-tune the full network.
+
+    Uses ``dp --pt train --finetune`` under the hood.  A GPU is strongly
+    recommended — this trains all parameters.
+    """
+    from deepmd.dpa_tools import DPAFineTuner
+
+    model = DPAFineTuner(
+        pretrained=model_path,
+        strategy="finetune",
+        property_name="gap",
+        task_dim=1,
+        intensive=True,
+        output_dir=str(HERE / "output_ft"),
+    )
+
+    print("finetune — fitting (dp --pt train) …")
+    model.fit(
+        train_data=str(TRAIN_DIR),
+        valid_data=str(TEST_DIR),
+        target_key="gap",
+    )
+
+    print("finetune — evaluating …")
+    metrics = model.evaluate(data=str(TEST_DIR))
+    _print_metrics(metrics, "finetune")
+
+
+# ---------------------------------------------------------------------------
+# Strategy 4 — mft (multi-task fine-tuning, GPU + aux data required)
+# ---------------------------------------------------------------------------
+
+
+def demo_mft(model_path: str) -> None:
+    """Multi-task fine-tuning: property head + auxiliary force-field head.
+
+    Requires auxiliary training data (e.g. SPICE2) via ``--aux-data``.
+    Prevents representation collapse on small property datasets.
+    """
+    from deepmd.dpa_tools import DPAFineTuner
+
+    model = DPAFineTuner(
+        pretrained=model_path,
+        strategy="mft",
+        property_name="gap",
+        aux_branch="MP_traj_v024_alldata_mixu",
+    )
+
+    print("mft — fitting …")
+    # NOTE: you must supply real aux_data; this is a placeholder.
+    model.fit(
+        train_data=str(TRAIN_DIR),
+        aux_data=str(HERE / "data" / "aux"),   # ← replace with your aux data path
+    )
+
+    print("mft — evaluating …")
+    metrics = model.evaluate(data=str(TEST_DIR))
+    _print_metrics(metrics, "mft")
+
+
+# ---------------------------------------------------------------------------
+# main
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="dp dpa fit",
+        description="Quickstart: fit a DPA model on QM9 HOMO-LUMO gap.",
+    )
+    parser.add_argument(
+        "--model",
+        default=None,
+        help="Path to DPA checkpoint (.pt).  Falls back to $DPA_MODEL_PATH.",
+    )
+    parser.add_argument(
+        "--strategy",
+        default="frozen_sklearn",
+        choices=["frozen_sklearn", "linear_probe", "finetune", "mft"],
+        help="Fine-tuning strategy (default: %(default)s).",
+    )
+    args = parser.parse_args()
+
+    model_path = _resolve_model(args)
+    _verify_data()
+    train_labels, _test_labels = _load_labels()
+
+    print(f"Model   : {model_path}")
+    print(f"Strategy: {args.strategy}")
+
+    if args.strategy == "frozen_sklearn":
+        demo_frozen_sklearn(model_path, train_labels)
+    elif args.strategy == "linear_probe":
+        demo_linear_probe(model_path)
+    elif args.strategy == "finetune":
+        demo_finetune(model_path)
+    elif args.strategy == "mft":
+        demo_mft(model_path)
+
+
 if __name__ == "__main__":
     main()