Skip to content

Commit e28d4dd

Browse files
zirenjinzirenjin
committed
feat(demo): demonstrate all four fine-tuning strategies
1 parent 892ab07 commit e28d4dd

1 file changed

Lines changed: 201 additions & 58 deletions

File tree

‎deepmd/dpa_tools/demo/fit_evaluate.py‎

Lines changed: 201 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,17 @@
11
#!/usr/bin/env python3
2-
"""Fit a frozen DPA descriptor + Ridge regressor on the quickstart demo data.
2+
"""Fit a pretrained DPA descriptor on the quickstart QM9 demo data.
33
4-
Requires the DPA-3.1-3M pretrained checkpoint. Provide it via ``--model`` or
5-
set the ``DPA_MODEL_PATH`` environment variable.
4+
All four fine-tuning strategies are demonstrated. The default
5+
(``frozen_sklearn``) runs on CPU in under 5 minutes and requires no GPU.
6+
``linear_probe``, ``finetune``, and ``mft`` use ``dp --pt train`` under the
7+
hood and need a GPU to finish in reasonable time.
68
79
Usage (from the demo directory)::
810
911
python fit_evaluate.py --model /path/to/DPA-3.1-3M.pt
12+
python fit_evaluate.py --model /path/to/DPA-3.1-3M.pt --strategy finetune
1013
11-
(or set the ``DPA_MODEL_PATH`` environment variable instead of ``--model``).
14+
Set ``DPA_MODEL_PATH`` instead of ``--model`` to avoid typing it every time.
1215
"""
1316

1417
from __future__ import annotations
@@ -29,90 +32,230 @@
2932
FROZEN_MODEL_PATH = HERE / "frozen_model.pth"
3033

3134

32-
def main() -> None:
33-
parser = argparse.ArgumentParser(
34-
prog="dp dpa fit",
35-
description="Quickstart: fit frozen DPA descriptor + Ridge on QM9 HOMO-LUMO gap.",
36-
)
37-
parser.add_argument(
38-
"--model",
39-
default=None,
40-
help="Path to DPA-3.1-3M.pt checkpoint. Falls back to $DPA_MODEL_PATH.",
41-
)
42-
args = parser.parse_args()
35+
# ---------------------------------------------------------------------------
36+
# helpers
37+
# ---------------------------------------------------------------------------
38+
4339

44-
# --- resolve model path ---
45-
model_path = args.model or os.environ.get("DPA_MODEL_PATH")
46-
if not model_path:
40+
def _resolve_model(args: argparse.Namespace) -> str:
41+
path = args.model or os.environ.get("DPA_MODEL_PATH")
42+
if not path:
4743
print(
48-
"error: DPA-3.1-3M checkpoint not specified.\n"
49-
" Provide it via --model or set the DPA_MODEL_PATH environment variable.\n"
44+
"error: DPA checkpoint not specified.\n"
45+
" Provide it via --model or set $DPA_MODEL_PATH.\n"
5046
" Example: python fit_evaluate.py --model /path/to/DPA-3.1-3M.pt",
5147
file=sys.stderr,
5248
)
5349
sys.exit(1)
54-
55-
if not Path(model_path).is_file():
56-
print(f"error: model file not found: {model_path}", file=sys.stderr)
50+
if not Path(path).is_file():
51+
print(f"error: model file not found: {path}", file=sys.stderr)
5752
sys.exit(1)
53+
return path
5854

59-
print(f"Model checkpoint: {model_path}")
6055

61-
# --- verify data ---
62-
if not TRAIN_DIR.is_dir():
63-
print(
64-
f"error: training data not found at {TRAIN_DIR}\n"
65-
" Run scripts/prepare_data.py first.",
66-
file=sys.stderr,
67-
)
68-
sys.exit(1)
69-
if not TEST_DIR.is_dir():
70-
print(
71-
f"error: test data not found at {TEST_DIR}\n"
72-
" Run scripts/prepare_data.py first.",
73-
file=sys.stderr,
74-
)
75-
sys.exit(1)
56+
def _verify_data() -> None:
57+
for name, d in [("train", TRAIN_DIR), ("test", TEST_DIR)]:
58+
if not d.is_dir():
59+
print(
60+
f"error: {name} data not found at {d}\n"
61+
" Run scripts/prepare_data.py first.",
62+
file=sys.stderr,
63+
)
64+
sys.exit(1)
65+
66+
67+
def _load_labels() -> tuple[np.ndarray, np.ndarray]:
68+
train = np.load(str(TRAIN_LABELS_PATH)).astype(np.float32)
69+
test = np.load(str(TEST_LABELS_PATH)).astype(np.float32)
70+
return train, test
71+
72+
73+
def _print_metrics(metrics, label: str = "") -> None:
74+
tag = f" [{label}]" if label else ""
75+
print()
76+
print("=" * 50)
77+
print(f"MAE{tag} : {metrics.mae:.4f} eV")
78+
print(f"R²{tag} : {metrics.r2:.4f}")
79+
print(f"RMSE{tag} : {metrics.rmse:.4f} eV")
80+
print(f"N{tag} : {metrics.predictions.shape[0]}")
81+
print("=" * 50)
7682

77-
# --- load labels ---
78-
train_labels = np.load(str(TRAIN_LABELS_PATH)).astype(np.float32)
79-
test_labels = np.load(str(TEST_LABELS_PATH)).astype(np.float32)
8083

81-
# --- build model ---
84+
# ---------------------------------------------------------------------------
85+
# Strategy 1 — frozen_sklearn (default, CPU)
86+
# ---------------------------------------------------------------------------
87+
88+
89+
def demo_frozen_sklearn(model_path: str, train_labels: np.ndarray) -> None:
90+
"""Freeze the DPA backbone, extract descriptors once, fit a sklearn Ridge.
91+
92+
Fastest iteration. No GPU, no ``dp train`` subprocess. The frozen
93+
bundle (``.pth``) is portable and can be loaded with ``DPAPredictor``.
94+
"""
8295
from deepmd.dpa_tools import DPAFineTuner
8396

8497
model = DPAFineTuner(
8598
pretrained=model_path,
8699
model_branch="Domains_Drug",
87-
pooling="mean",
88-
predictor="linear",
100+
strategy="frozen_sklearn",
101+
predictor="linear", # "linear" (Ridge) | "rf" | "mlp"
102+
pooling="mean", # "mean" | "sum" | "mean+std" | "mean+std+max+min"
89103
seed=42,
90104
)
91105

92-
# --- fit ---
93-
print("Fitting …")
106+
print("frozen_sklearn — fitting …")
94107
model.fit(
95108
train_data=str(TRAIN_DIR),
96109
labels=train_labels,
97110
target_key="gap",
98111
)
99112

100-
# --- evaluate ---
101-
print("Evaluating …")
113+
print("frozen_sklearn — evaluating …")
102114
metrics = model.evaluate(data=str(TEST_DIR))
115+
_print_metrics(metrics, "frozen_sklearn")
103116

104-
print()
105-
print("=" * 50)
106-
print(f"MAE : {metrics.mae:.4f} eV")
107-
print(f"R² : {metrics.r2:.4f}")
108-
print(f"RMSE : {metrics.rmse:.4f} eV")
109-
print(f"N : {metrics.predictions.shape[0]}")
110-
print("=" * 50)
111-
112-
# --- freeze ---
113117
out = model.freeze(str(FROZEN_MODEL_PATH))
114118
print(f"Frozen model → {out}")
115119

116120

121+
# ---------------------------------------------------------------------------
122+
# Strategy 2 — linear_probe (GPU recommended)
123+
# ---------------------------------------------------------------------------
124+
125+
126+
def demo_linear_probe(model_path: str) -> None:
127+
"""Freeze the DPA backbone, train only a neural property fitting net.
128+
129+
Uses ``dp --pt train --finetune`` under the hood. A GPU is recommended.
130+
"""
131+
from deepmd.dpa_tools import DPAFineTuner
132+
133+
model = DPAFineTuner(
134+
pretrained=model_path,
135+
strategy="linear_probe",
136+
property_name="gap",
137+
task_dim=1,
138+
intensive=True,
139+
output_dir=str(HERE / "output_lp"),
140+
)
141+
142+
print("linear_probe — fitting (dp --pt train) …")
143+
model.fit(
144+
train_data=str(TRAIN_DIR),
145+
valid_data=str(TEST_DIR),
146+
target_key="gap",
147+
)
148+
149+
print("linear_probe — evaluating …")
150+
metrics = model.evaluate(data=str(TEST_DIR))
151+
_print_metrics(metrics, "linear_probe")
152+
153+
154+
# ---------------------------------------------------------------------------
155+
# Strategy 3 — finetune (GPU recommended)
156+
# ---------------------------------------------------------------------------
157+
158+
159+
def demo_finetune(model_path: str) -> None:
160+
"""Load the pretrained backbone and fine-tune the full network.
161+
162+
Uses ``dp --pt train --finetune`` under the hood. A GPU is strongly
163+
recommended — this trains all parameters.
164+
"""
165+
from deepmd.dpa_tools import DPAFineTuner
166+
167+
model = DPAFineTuner(
168+
pretrained=model_path,
169+
strategy="finetune",
170+
property_name="gap",
171+
task_dim=1,
172+
intensive=True,
173+
output_dir=str(HERE / "output_ft"),
174+
)
175+
176+
print("finetune — fitting (dp --pt train) …")
177+
model.fit(
178+
train_data=str(TRAIN_DIR),
179+
valid_data=str(TEST_DIR),
180+
target_key="gap",
181+
)
182+
183+
print("finetune — evaluating …")
184+
metrics = model.evaluate(data=str(TEST_DIR))
185+
_print_metrics(metrics, "finetune")
186+
187+
188+
# ---------------------------------------------------------------------------
189+
# Strategy 4 — mft (multi-task fine-tuning, GPU + aux data required)
190+
# ---------------------------------------------------------------------------
191+
192+
193+
def demo_mft(model_path: str) -> None:
194+
"""Multi-task fine-tuning: property head + auxiliary force-field head.
195+
196+
Requires auxiliary training data (e.g. SPICE2) via ``--aux-data``.
197+
Prevents representation collapse on small property datasets.
198+
"""
199+
from deepmd.dpa_tools import DPAFineTuner
200+
201+
model = DPAFineTuner(
202+
pretrained=model_path,
203+
strategy="mft",
204+
property_name="gap",
205+
aux_branch="MP_traj_v024_alldata_mixu",
206+
)
207+
208+
print("mft — fitting …")
209+
# NOTE: you must supply real aux_data; this is a placeholder.
210+
model.fit(
211+
train_data=str(TRAIN_DIR),
212+
aux_data=str(HERE / "data" / "aux"), # ← replace with your aux data path
213+
)
214+
215+
print("mft — evaluating …")
216+
metrics = model.evaluate(data=str(TEST_DIR))
217+
_print_metrics(metrics, "mft")
218+
219+
220+
# ---------------------------------------------------------------------------
221+
# main
222+
# ---------------------------------------------------------------------------
223+
224+
225+
def main() -> None:
226+
parser = argparse.ArgumentParser(
227+
prog="dp dpa fit",
228+
description="Quickstart: fit a DPA model on QM9 HOMO-LUMO gap.",
229+
)
230+
parser.add_argument(
231+
"--model",
232+
default=None,
233+
help="Path to DPA checkpoint (.pt). Falls back to $DPA_MODEL_PATH.",
234+
)
235+
parser.add_argument(
236+
"--strategy",
237+
default="frozen_sklearn",
238+
choices=["frozen_sklearn", "linear_probe", "finetune", "mft"],
239+
help="Fine-tuning strategy (default: %(default)s).",
240+
)
241+
args = parser.parse_args()
242+
243+
model_path = _resolve_model(args)
244+
_verify_data()
245+
train_labels, _test_labels = _load_labels()
246+
247+
print(f"Model : {model_path}")
248+
print(f"Strategy: {args.strategy}")
249+
250+
if args.strategy == "frozen_sklearn":
251+
demo_frozen_sklearn(model_path, train_labels)
252+
elif args.strategy == "linear_probe":
253+
demo_linear_probe(model_path)
254+
elif args.strategy == "finetune":
255+
demo_finetune(model_path)
256+
elif args.strategy == "mft":
257+
demo_mft(model_path)
258+
259+
117260
if __name__ == "__main__":
118261
main()

0 commit comments

Comments
 (0)