diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index afde4abdec..756d427a2c 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -621,6 +621,22 @@ def _eval_model( dtype=torch.float64, device=DEVICE, ) + elif self._is_pt2 and self.get_dim_fparam() > 0: + # .pt2 models are compiled with fparam as a required input. + # When the user omits fparam, fill with default values from metadata. + default_fp = self.metadata.get("default_fparam") + if default_fp is not None: + fparam_t = ( + torch.tensor(default_fp, dtype=torch.float64, device=DEVICE) + .unsqueeze(0) + .expand(nframes, -1) + .contiguous() + ) + else: + raise ValueError( + f"fparam is required for this model (dim_fparam={self.get_dim_fparam()}) " + "but was not provided, and no default_fparam is stored in the model." + ) else: fparam_t = None diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index b7e7c3edc7..35c32bc562 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -195,6 +195,7 @@ def _collect_metadata(model: torch.nn.Module) -> dict: "mixed_types": model.mixed_types(), "sel_type": model.get_sel_type(), "has_default_fparam": model.has_default_fparam(), + "default_fparam": model.get_default_fparam(), "fitting_output_defs": fitting_output_defs, } diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index e086334910..0d42324d24 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -202,6 +202,7 @@ class DeepPotPTExpt : public DeepPotBackend { int daparam; bool aparam_nall; bool has_default_fparam_; + std::vector default_fparam_; double rcut; int gpu_id; bool gpu_enabled; diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 267db796d8..5770b80264 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -636,6 +636,13 @@ void DeepPotPTExpt::init(const std::string& model, } else { has_default_fparam_ = false; } + default_fparam_.clear(); + if (has_default_fparam_ && metadata.obj_val.count("default_fparam") && + metadata["default_fparam"].type == JsonValue::Array) { + for (const auto& v : metadata["default_fparam"].as_array()) { + default_fparam_.push_back(v.as_double()); + } + } type_map.clear(); for (const auto& v : metadata["type_map"].as_array()) { @@ -818,6 +825,18 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, valuetype_options) .to(torch::kFloat64) .to(device); + } else if (!default_fparam_.empty()) { + fparam_tensor = + torch::from_blob(const_cast(default_fparam_.data()), + {1, static_cast(default_fparam_.size())}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); + } else if (dfparam > 0) { + throw deepmd::deepmd_exception( + "fparam is required for this model (dim_fparam=" + + std::to_string(dfparam) + + ") but was not provided, and no default_fparam is stored."); } else { fparam_tensor = torch::zeros({0}, options).to(device); } @@ -982,6 +1001,15 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, min_z = std::min(min_z, coord_d[ii * 3 + 2]); max_z = std::max(max_z, coord_d[ii * 3 + 2]); } + // Shift coords so minimum is at rcut (ensures all atoms are in [0, L)) + double shift_x = rcut - min_x; + double shift_y = rcut - min_y; + double shift_z = rcut - min_z; + for (int ii = 0; ii < natoms; ++ii) { + coord_d[ii * 3 + 0] += shift_x; + coord_d[ii * 3 + 1] += shift_y; + coord_d[ii * 3 + 2] += shift_z; + } box_d.resize(9, 0.0); box_d[0] = (max_x - min_x) + 2.0 * rcut; box_d[4] = (max_y - min_y) + 2.0 * rcut; @@ -1052,6 +1080,18 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, valuetype_options) .to(torch::kFloat64) .to(device); + } else if (!default_fparam_.empty()) { + fparam_tensor = + torch::from_blob(const_cast(default_fparam_.data()), + {1, static_cast(default_fparam_.size())}, + torch::TensorOptions().dtype(torch::kFloat64)) + .clone() + .to(device); + } else if (dfparam > 0) { + throw deepmd::deepmd_exception( + "fparam is required for this model (dim_fparam=" + + std::to_string(dfparam) + + ") but was not provided, and no default_fparam is stored."); } else { fparam_tensor = torch::zeros({0}, options).to(device); } diff --git a/source/api_cc/tests/test_deeppot_a_fparam_aparam_ptexpt.cc b/source/api_cc/tests/test_deeppot_a_fparam_aparam_ptexpt.cc index 829a821e9e..517620d72f 100644 --- a/source/api_cc/tests/test_deeppot_a_fparam_aparam_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_a_fparam_aparam_ptexpt.cc @@ -370,3 +370,89 @@ TYPED_TEST(TestInferDeepPotDefaultFParamPtExpt, has_default_fparam) { EXPECT_EQ(dp.dim_fparam(), 1); EXPECT_TRUE(dp.has_default_fparam()); } + +// Eval without fparam should produce same result as eval with explicit default +TYPED_TEST(TestInferDeepPotDefaultFParamPtExpt, eval_default_vs_explicit) { + using VALUETYPE = TypeParam; + deepmd::DeepPot& dp = this->dp; + + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 0, 0, 0, 0, 0}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + // The default fparam value from gen_fparam_aparam.py + std::vector explicit_fparam = {0.25852028}; + std::vector aparam = {0.25852028, 0.25852028, 0.25852028, + 0.25852028, 0.25852028, 0.25852028}; + + // Eval with explicit fparam + double e_explicit; + std::vector f_explicit, v_explicit; + dp.compute(e_explicit, f_explicit, v_explicit, coord, atype, box, + explicit_fparam, aparam); + + // Eval without fparam (should use default) + double e_default; + std::vector f_default, v_default; + std::vector empty_fparam; + dp.compute(e_default, f_default, v_default, coord, atype, box, empty_fparam, + aparam); + + EXPECT_LT(fabs(e_explicit - e_default), EPSILON); + for (size_t ii = 0; ii < f_explicit.size(); ++ii) { + EXPECT_LT(fabs(f_explicit[ii] - f_default[ii]), EPSILON); + } + for (size_t ii = 0; ii < v_explicit.size(); ++ii) { + EXPECT_LT(fabs(v_explicit[ii] - v_default[ii]), EPSILON); + } +} + +// Same test but with external nlist (LAMMPS path) +TYPED_TEST(TestInferDeepPotDefaultFParamPtExpt, eval_default_vs_explicit_lmp) { + using VALUETYPE = TypeParam; + deepmd::DeepPot& dp = this->dp; + + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 0, 0, 0, 0, 0}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + std::vector explicit_fparam = {0.25852028}; + std::vector aparam = {0.25852028, 0.25852028, 0.25852028, + 0.25852028, 0.25852028, 0.25852028}; + + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + // Eval with explicit fparam (external nlist) + double e_explicit; + std::vector f_explicit, v_explicit; + dp.compute(e_explicit, f_explicit, v_explicit, coord_cpy, atype_cpy, box, + nall - nloc, inlist, 0, explicit_fparam, aparam); + + // Eval without fparam (external nlist, should use default) + double e_default; + std::vector f_default, v_default; + std::vector empty_fparam; + dp.compute(e_default, f_default, v_default, coord_cpy, atype_cpy, box, + nall - nloc, inlist, 0, empty_fparam, aparam); + + EXPECT_LT(fabs(e_explicit - e_default), EPSILON); + for (size_t ii = 0; ii < f_explicit.size(); ++ii) { + EXPECT_LT(fabs(f_explicit[ii] - f_default[ii]), EPSILON); + } + for (size_t ii = 0; ii < v_explicit.size(); ++ii) { + EXPECT_LT(fabs(v_explicit[ii] - v_default[ii]), EPSILON); + } +} diff --git a/source/api_cc/tests/test_deeppot_ptexpt.cc b/source/api_cc/tests/test_deeppot_ptexpt.cc index 6a1cabefba..2cda44bb19 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_ptexpt.cc @@ -19,6 +19,10 @@ class TestInferDeepPotAPtExpt : public ::testing::Test { std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + // Alternative coords for multi-frame tests (must give different energy) + std::vector coord_alt = {10.06, 5.71, 11.16, 9.07, 1.22, 12.68, + 9.89, 10.22, 1.67, 5.86, 4.82, 12.05, + 8.37, 10.70, 5.76, 2.95, 7.21, 0.83}; std::vector atype = {0, 1, 1, 0, 1, 1}; std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; // Same reference values as test_deeppot_pt.cc (model converted from .pth) @@ -414,6 +418,9 @@ class TestInferDeepPotAPtExptNoPbc : public ::testing::Test { std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector coord_alt = {10.06, 5.71, 11.16, 9.07, 1.22, 12.68, + 9.89, 10.22, 1.67, 5.86, 4.82, 12.05, + 8.37, 10.70, 5.76, 2.95, 7.21, 0.83}; std::vector atype = {0, 1, 1, 0, 1, 1}; std::vector box = {}; // Same reference values as TestInferDeepPotAPtNoPbc in test_deeppot_pt.cc @@ -500,3 +507,259 @@ TYPED_TEST(TestInferDeepPotAPtExptNoPbc, cpu_build_nlist) { EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); } } + +TYPED_TEST(TestInferDeepPotAPtExptNoPbc, cpu_build_nlist_atomic) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_e = this->expected_e; + std::vector& expected_f = this->expected_f; + std::vector& expected_v = this->expected_v; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + double ener; + std::vector force, virial, atom_ener, atom_vir; + dp.compute(ener, force, virial, atom_ener, atom_vir, coord, atype, box); + + EXPECT_EQ(force.size(), natoms * 3); + EXPECT_EQ(virial.size(), 9); + EXPECT_EQ(atom_ener.size(), natoms); + EXPECT_EQ(atom_vir.size(), natoms * 9); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 3 * 3; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + for (int ii = 0; ii < natoms; ++ii) { + EXPECT_LT(fabs(atom_ener[ii] - expected_e[ii]), EPSILON); + } + for (int ii = 0; ii < natoms * 9; ++ii) { + EXPECT_LT(fabs(atom_vir[ii] - expected_v[ii]), EPSILON); + } +} + +// Multi-frame PBC test via compute_mixed_type +TYPED_TEST(TestInferDeepPotAPtExpt, cpu_build_nlist_nframes) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& coord_alt = this->coord_alt; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + + int nframes = 2; + // Frame 0: original coords. Frame 1: alternative coords (coord_alt). + std::vector coord_2f(coord); + coord_2f.insert(coord_2f.end(), coord_alt.begin(), coord_alt.end()); + std::vector atype_2f(atype); + atype_2f.insert(atype_2f.end(), atype.begin(), atype.end()); + std::vector box_2f(box); + box_2f.insert(box_2f.end(), box.begin(), box.end()); + + std::vector ener; + std::vector force, virial; + dp.compute_mixed_type(ener, force, virial, nframes, coord_2f, atype_2f, + box_2f); + + EXPECT_EQ(ener.size(), nframes); + EXPECT_EQ(force.size(), nframes * natoms * 3); + EXPECT_EQ(virial.size(), nframes * 9); + + // Frame 0 should match reference + EXPECT_LT(fabs(ener[0] - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + // Frame 1 should be different (perturbed coords) + EXPECT_GT(fabs(ener[1] - ener[0]), 1e-10); +} + +// Multi-frame NoPBC test via compute_mixed_type +TYPED_TEST(TestInferDeepPotAPtExptNoPbc, cpu_build_nlist_nframes) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& coord_alt = this->coord_alt; + std::vector& box = this->box; // empty + std::vector& expected_f = this->expected_f; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e; + std::vector& expected_tot_v = this->expected_tot_v; + deepmd::DeepPot& dp = this->dp; + + int nframes = 2; + std::vector coord_2f(coord); + coord_2f.insert(coord_2f.end(), coord_alt.begin(), coord_alt.end()); + std::vector atype_2f(atype); + atype_2f.insert(atype_2f.end(), atype.begin(), atype.end()); + + std::vector ener; + std::vector force, virial; + dp.compute_mixed_type(ener, force, virial, nframes, coord_2f, atype_2f, box); + + EXPECT_EQ(ener.size(), nframes); + EXPECT_EQ(force.size(), nframes * natoms * 3); + EXPECT_EQ(virial.size(), nframes * 9); + + // Frame 0 should match reference + EXPECT_LT(fabs(ener[0] - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + // Frame 1 should be different (perturbed coords) + EXPECT_GT(fabs(ener[1] - ener[0]), 1e-10); +} + +// ========== Parser / metadata coverage tests ========== + +TEST(TestDeepPotPTExptParser, load_nonexistent_file) { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd::DeepPot dp; + EXPECT_THROW(dp.init("nonexistent_model.pt2"), deepmd::deepmd_exception); +} + +TEST(TestDeepPotPTExptParser, load_invalid_zip) { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::string tmpfile = "test_invalid.pt2"; + { + std::ofstream ofs(tmpfile, std::ios::binary); + ASSERT_TRUE(ofs.is_open()) << "Failed to create temp file"; + ofs << "not a zip file at all"; + } + deepmd::DeepPot dp; + EXPECT_THROW(dp.init(tmpfile), deepmd::deepmd_exception); + std::remove(tmpfile.c_str()); +} + +TEST(TestDeepPotPTExptParser, load_tiny_file) { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::string tmpfile = "test_tiny.pt2"; + { + std::ofstream ofs(tmpfile, std::ios::binary); + ASSERT_TRUE(ofs.is_open()) << "Failed to create temp file"; + ofs << "abc"; + } + deepmd::DeepPot dp; + EXPECT_THROW(dp.init(tmpfile), deepmd::deepmd_exception); + std::remove(tmpfile.c_str()); +} + +// Metadata accessor tests — exercise JSON parser on a real model +template +class TestDeepPotPTExptMetadata : public ::testing::Test { + protected: + deepmd::DeepPot dp; + void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + dp.init("../../tests/infer/deeppot_sea.pt2"); + }; + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestDeepPotPTExptMetadata, ValueTypes); + +TYPED_TEST(TestDeepPotPTExptMetadata, type_map) { + std::string type_map; + this->dp.get_type_map(type_map); + EXPECT_NE(type_map.find("O"), std::string::npos); + EXPECT_NE(type_map.find("H"), std::string::npos); +} + +TYPED_TEST(TestDeepPotPTExptMetadata, cutoff) { + EXPECT_GT(this->dp.cutoff(), 0.0); +} + +TYPED_TEST(TestDeepPotPTExptMetadata, ntypes) { + EXPECT_EQ(this->dp.numb_types(), 2); +} + +TYPED_TEST(TestDeepPotPTExptMetadata, dim_fparam_zero) { + EXPECT_EQ(this->dp.dim_fparam(), 0); +} + +TYPED_TEST(TestDeepPotPTExptMetadata, dim_aparam_zero) { + EXPECT_EQ(this->dp.dim_aparam(), 0); +} + +TYPED_TEST(TestDeepPotPTExptMetadata, no_default_fparam) { + EXPECT_FALSE(this->dp.has_default_fparam()); +} + +// JSON parser type-coverage via fparam model +template +class TestDeepPotPTExptJsonTypes : public ::testing::Test { + protected: + deepmd::DeepPot dp; + void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + dp.init("../../tests/infer/fparam_aparam.pt2"); + }; + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestDeepPotPTExptJsonTypes, ValueTypes); + +TYPED_TEST(TestDeepPotPTExptJsonTypes, integer_fields) { + EXPECT_EQ(this->dp.dim_fparam(), 1); + EXPECT_EQ(this->dp.dim_aparam(), 1); +} + +TYPED_TEST(TestDeepPotPTExptJsonTypes, boolean_field) { + EXPECT_FALSE(this->dp.has_default_fparam()); +} + +TYPED_TEST(TestDeepPotPTExptJsonTypes, string_array) { + std::string type_map; + this->dp.get_type_map(type_map); + EXPECT_FALSE(type_map.empty()); +} + +TYPED_TEST(TestDeepPotPTExptJsonTypes, float_field) { + EXPECT_GT(this->dp.cutoff(), 0.0); + EXPECT_LT(this->dp.cutoff(), 100.0); +} + +// Default fparam model — tests JSON parsing of boolean true + float array +template +class TestDeepPotPTExptJsonDefaults : public ::testing::Test { + protected: + deepmd::DeepPot dp; + void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + dp.init("../../tests/infer/fparam_aparam_default.pt2"); + }; + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestDeepPotPTExptJsonDefaults, ValueTypes); + +TYPED_TEST(TestDeepPotPTExptJsonDefaults, boolean_true) { + EXPECT_TRUE(this->dp.has_default_fparam()); +} + +TYPED_TEST(TestDeepPotPTExptJsonDefaults, default_fparam_parsed) { + EXPECT_EQ(this->dp.dim_fparam(), 1); + EXPECT_TRUE(this->dp.has_default_fparam()); +} diff --git a/source/tests/pt_expt/model/test_model_compression.py b/source/tests/pt_expt/model/test_model_compression.py index 619bd10981..fb7681a8bb 100644 --- a/source/tests/pt_expt/model/test_model_compression.py +++ b/source/tests/pt_expt/model/test_model_compression.py @@ -290,6 +290,83 @@ def test_compress_cli_entry_point(self) -> None: if os.path.exists(compressed_path): os.unlink(compressed_path) + def test_freeze_compress_eval_pt2(self) -> None: + """Test pipeline: build → freeze (.pte) → compress → output (.pt2) → eval. + + Verifies that compressed models can be exported to .pt2 format and + produce consistent results when loaded via DeepPot. + """ + from deepmd.infer import ( + DeepPot, + ) + from deepmd.pt_expt.entrypoints.compress import ( + enable_compression as compress_entry, + ) + + # 1. Build and freeze to .pte + md = self._make_model() + md.min_nbor_dist = 0.5 + md.eval() + ret_frozen = self._eval_model(md) + + model_data = {"model": md.serialize(), "min_nbor_dist": 0.5} + with tempfile.NamedTemporaryFile(suffix=".pte", delete=False) as f: + frozen_path = f.name + with tempfile.NamedTemporaryFile(suffix=".pt2", delete=False) as f: + compressed_pt2_path = f.name + try: + deserialize_to_file(frozen_path, model_data) + + # 2. Compress with .pt2 output + compress_entry( + input_file=frozen_path, + output=compressed_pt2_path, + stride=0.01, + extrapolate=5, + check_frequency=-1, + ) + self.assertTrue(os.path.exists(compressed_pt2_path)) + + # 3. Load .pt2 and verify metadata has compression state + compressed_data = serialize_from_file(compressed_pt2_path) + descrpt_data = compressed_data["model"]["descriptor"] + self.assertIn("compress", descrpt_data) + + # 4. Eval via DeepPot and verify consistency + dp = DeepPot(compressed_pt2_path) + coord = self.coord.detach().cpu().numpy().reshape(-1) + box = self.cell.reshape(9).detach().cpu().numpy() + atype = self.atype[0].detach().cpu().numpy().tolist() + e_pt2, f_pt2, v_pt2 = dp.eval(coord, box, atype) + np.testing.assert_allclose( + ret_frozen["energy"], e_pt2, atol=1e-7, err_msg="energy" + ) + np.testing.assert_allclose( + ret_frozen["force"], f_pt2, atol=1e-7, err_msg="force" + ) + np.testing.assert_allclose( + ret_frozen["virial"], v_pt2, atol=1e-7, err_msg="virial" + ) + finally: + os.unlink(frozen_path) + if os.path.exists(compressed_pt2_path): + os.unlink(compressed_pt2_path) + + def test_min_nbor_dist_roundtrip_pt2(self) -> None: + """Test that min_nbor_dist survives freeze → load round-trip via .pt2.""" + md = self._make_model() + md.min_nbor_dist = 0.5 + + model_data = {"model": md.serialize(), "min_nbor_dist": 0.5} + with tempfile.NamedTemporaryFile(suffix=".pt2", delete=False) as f: + pt2_path = f.name + try: + deserialize_to_file(pt2_path, model_data) + loaded_data = serialize_from_file(pt2_path) + self.assertAlmostEqual(loaded_data["min_nbor_dist"], 0.5) + finally: + os.unlink(pt2_path) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_change_bias.py b/source/tests/pt_expt/test_change_bias.py index 16974cc653..65cb785500 100644 --- a/source/tests/pt_expt/test_change_bias.py +++ b/source/tests/pt_expt/test_change_bias.py @@ -153,6 +153,16 @@ def setUpClass(cls) -> None: # Record original bias cls.original_bias = to_numpy(trainer.wrapper.model.get_out_bias()) + # Freeze once to .pte and .pt2 for reuse across tests + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + + cls.pte_path = os.path.join(cls.tmpdir, "shared.pte") + cls.pt2_path = os.path.join(cls.tmpdir, "shared.pt2") + freeze(model=cls.model_path, output=cls.pte_path) + freeze(model=cls.model_path, output=cls.pt2_path) + @classmethod def tearDownClass(cls) -> None: os.chdir(cls.old_cwd) @@ -219,9 +229,6 @@ def test_change_bias_with_user_defined(self) -> None: np.testing.assert_allclose(updated_bias, expected_bias) def test_change_bias_frozen_pte(self) -> None: - from deepmd.pt_expt.entrypoints.main import ( - freeze, - ) from deepmd.pt_expt.model.model import ( BaseModel, ) @@ -229,19 +236,15 @@ def test_change_bias_frozen_pte(self) -> None: serialize_from_file, ) - # Freeze the checkpoint - pte_path = os.path.join(self.tmpdir, "frozen.pte") - freeze(model=self.model_path, output=pte_path) - - # Get original bias - original_data = serialize_from_file(pte_path) + # Get original bias from shared frozen model + original_data = serialize_from_file(self.pte_path) original_model = BaseModel.deserialize(original_data["model"]) original_bias = to_numpy(original_model.get_out_bias()) # Run change-bias on the frozen model output_pte = os.path.join(self.tmpdir, "frozen_updated.pte") run_dp( - f"dp --pt-expt change-bias {pte_path} " + f"dp --pt-expt change-bias {self.pte_path} " f"-s {self.data_file[0]} -o {output_pte}" ) @@ -250,12 +253,85 @@ def test_change_bias_frozen_pte(self) -> None: updated_model = BaseModel.deserialize(updated_data["model"]) updated_bias = to_numpy(updated_model.get_out_bias()) - # Bias should have changed self.assertFalse( np.allclose(original_bias, updated_bias), "Bias should have changed after change-bias on frozen model", ) + def test_change_bias_frozen_pt2(self) -> None: + """Change-bias on a .pt2 frozen model.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + serialize_from_file, + ) + + original_data = serialize_from_file(self.pt2_path) + original_model = BaseModel.deserialize(original_data["model"]) + original_bias = to_numpy(original_model.get_out_bias()) + + output_pt2 = os.path.join(self.tmpdir, "frozen_updated.pt2") + run_dp( + f"dp --pt-expt change-bias {self.pt2_path} " + f"-s {self.data_file[0]} -o {output_pt2}" + ) + + updated_data = serialize_from_file(output_pt2) + updated_model = BaseModel.deserialize(updated_data["model"]) + updated_bias = to_numpy(updated_model.get_out_bias()) + + self.assertFalse( + np.allclose(original_bias, updated_bias), + "Bias should have changed after change-bias on .pt2 model", + ) + + def test_change_bias_frozen_pt2_user_defined(self) -> None: + """Change-bias with user-defined values on a .pt2 model.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + serialize_from_file, + ) + + output_pt2 = os.path.join(self.tmpdir, "frozen_ud_updated.pt2") + run_dp(f"dp --pt-expt change-bias {self.pt2_path} -b 1.0 2.0 -o {output_pt2}") + + updated_data = serialize_from_file(output_pt2) + updated_model = BaseModel.deserialize(updated_data["model"]) + updated_bias = to_numpy(updated_model.get_out_bias()) + + expected = np.array([1.0, 2.0]).reshape(updated_bias.shape) + np.testing.assert_allclose(updated_bias, expected, atol=1e-10) + + def test_change_bias_pt2_pte_consistency(self) -> None: + """Change-bias on .pte and .pt2 should produce same bias values.""" + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.serialization import ( + serialize_from_file, + ) + + output_pte = os.path.join(self.tmpdir, "cons_updated.pte") + output_pt2 = os.path.join(self.tmpdir, "cons_updated.pt2") + run_dp( + f"dp --pt-expt change-bias {self.pte_path} " + f"-s {self.data_file[0]} -o {output_pte}" + ) + run_dp( + f"dp --pt-expt change-bias {self.pt2_path} " + f"-s {self.data_file[0]} -o {output_pt2}" + ) + + pte_data = serialize_from_file(output_pte) + pt2_data = serialize_from_file(output_pt2) + pte_bias = to_numpy(BaseModel.deserialize(pte_data["model"]).get_out_bias()) + pt2_bias = to_numpy(BaseModel.deserialize(pt2_data["model"]).get_out_bias()) + + np.testing.assert_allclose(pte_bias, pt2_bias, atol=1e-10) + class TestChangeBiasFittingStats(unittest.TestCase): """Test that model_change_out_bias recomputes fitting stats for set-by-statistic.""" diff --git a/source/tests/pt_expt/test_dp_freeze.py b/source/tests/pt_expt/test_dp_freeze.py index f1c78eda16..bdbc206a6e 100644 --- a/source/tests/pt_expt/test_dp_freeze.py +++ b/source/tests/pt_expt/test_dp_freeze.py @@ -8,6 +8,7 @@ deepcopy, ) +import numpy as np import torch from deepmd.pt_expt.entrypoints.main import ( @@ -25,28 +26,51 @@ "type_map": ["O", "H", "B"], "descriptor": { "type": "se_e2_a", - "sel": [46, 92, 4], + "sel": [6, 12, 4], "rcut_smth": 0.50, "rcut": 4.00, - "neuron": [25, 50, 100], + "neuron": [8, 16], "resnet_dt": False, - "axis_neuron": 16, + "axis_neuron": 4, "seed": 1, }, "fitting_net": { - "neuron": [24, 24, 24], + "neuron": [8, 8], "resnet_dt": True, "seed": 1, }, "data_stat_nbatch": 20, } +# Shared test coordinates +_coord = np.array( + [0.0, 0.0, 0.1, 1.0, 0.3, 0.2, 0.1, 1.9, 3.4], + dtype=np.float64, +) +_box = np.array( + [5.0, 0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 5.0], + dtype=np.float64, +) +_atype = [0, 1, 2] + +# Coordinates with negative values for NoPBC tests +_coord_neg = np.array( + [-1.0, -2.0, 0.5, 1.0, 0.3, -0.2, 0.1, -1.9, 3.4], + dtype=np.float64, +) + class TestDPFreezePtExpt(unittest.TestCase): """Test dp freeze for the pt_expt backend.""" @classmethod def setUpClass(cls) -> None: + from .conftest import ( + _pop_device_contexts, + ) + + _pop_device_contexts() + cls.tmpdir = tempfile.mkdtemp() # Build a model and save a fake checkpoint @@ -57,15 +81,19 @@ def setUpClass(cls) -> None: cls.ckpt_file = os.path.join(cls.tmpdir, "model.pt") torch.save({"model": state_dict}, cls.ckpt_file) + # Freeze once to .pte and .pt2 for reuse across tests + cls.pte_path = os.path.join(cls.tmpdir, "shared.pte") + cls.pt2_path = os.path.join(cls.tmpdir, "shared.pt2") + freeze(model=cls.ckpt_file, output=cls.pte_path) + freeze(model=cls.ckpt_file, output=cls.pt2_path) + @classmethod def tearDownClass(cls) -> None: shutil.rmtree(cls.tmpdir) def test_freeze_pte(self) -> None: """Freeze to .pte and verify the file is created.""" - output = os.path.join(self.tmpdir, "frozen_model.pte") - freeze(model=self.ckpt_file, output=output) - self.assertTrue(os.path.exists(output)) + self.assertTrue(os.path.exists(self.pte_path)) def test_freeze_main_dispatcher(self) -> None: """Test main() CLI dispatcher with freeze command.""" @@ -96,6 +124,114 @@ def test_freeze_default_suffix(self) -> None: expected = os.path.join(self.tmpdir, "frozen_default_suffix.pte") self.assertTrue(os.path.exists(expected)) + def test_freeze_pt2(self) -> None: + """Freeze to .pt2 (AOTInductor) and verify the file is loadable.""" + self.assertTrue(os.path.exists(self.pt2_path)) + + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(self.pt2_path) + self.assertEqual(dp.get_type_map(), ["O", "H", "B"]) + rcut = dp.get_rcut() + self.assertGreater(rcut, 0.0) + + # Quick smoke-test eval + e, f, v = dp.eval(_coord, _box, _atype) + self.assertEqual(e.shape, (1, 1)) + self.assertEqual(f.shape, (1, 3, 3)) + self.assertEqual(v.shape, (1, 9)) + + def test_freeze_pt2_eval_consistency(self) -> None: + """Verify .pte and .pt2 produce identical results.""" + from deepmd.infer import ( + DeepPot, + ) + + dp_pte = DeepPot(self.pte_path) + dp_pt2 = DeepPot(self.pt2_path) + + e_pte, f_pte, v_pte = dp_pte.eval(_coord, _box, _atype) + e_pt2, f_pt2, v_pt2 = dp_pt2.eval(_coord, _box, _atype) + + self.assertTrue(np.isfinite(e_pte).all()) + self.assertTrue(np.isfinite(f_pte).all()) + self.assertTrue(np.isfinite(v_pte).all()) + np.testing.assert_allclose(e_pte, e_pt2, rtol=0, atol=1e-10) + np.testing.assert_allclose(f_pte, f_pt2, rtol=0, atol=1e-10) + np.testing.assert_allclose(v_pte, v_pt2, rtol=0, atol=1e-10) + + def test_freeze_pt2_nopbc_negative_coords(self) -> None: + """Verify .pt2 NoPBC works with negative coordinates.""" + from deepmd.infer import ( + DeepPot, + ) + + dp_pte = DeepPot(self.pte_path) + dp_pt2 = DeepPot(self.pt2_path) + + e_pte, f_pte, v_pte = dp_pte.eval(_coord_neg, None, _atype) + e_pt2, f_pt2, v_pt2 = dp_pt2.eval(_coord_neg, None, _atype) + + self.assertTrue(np.isfinite(e_pte).all()) + self.assertTrue(np.isfinite(f_pte).all()) + self.assertTrue(np.isfinite(v_pte).all()) + np.testing.assert_allclose(e_pte, e_pt2, rtol=0, atol=1e-10) + np.testing.assert_allclose(f_pte, f_pt2, rtol=0, atol=1e-10) + np.testing.assert_allclose(v_pte, v_pt2, rtol=0, atol=1e-10) + + +class TestDPFreezePt2DefaultFparam(unittest.TestCase): + """Test .pt2 with default fparam — eval without providing fparam.""" + + @classmethod + def setUpClass(cls) -> None: + from .conftest import ( + _pop_device_contexts, + ) + + _pop_device_contexts() + + cls.tmpdir = tempfile.mkdtemp() + + model_params = deepcopy(model_se_e2_a) + model_params["fitting_net"]["numb_fparam"] = 1 + model_params["fitting_net"]["default_fparam"] = [0.5] + model = get_model(model_params) + wrapper = ModelWrapper(model, model_params=model_params) + state_dict = wrapper.state_dict() + cls.ckpt_file = os.path.join(cls.tmpdir, "model_dfp.pt") + torch.save({"model": state_dict}, cls.ckpt_file) + + # Freeze once for reuse + cls.pt2_path = os.path.join(cls.tmpdir, "dfp.pt2") + freeze(model=cls.ckpt_file, output=cls.pt2_path) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls.tmpdir) + + def test_pt2_eval_default_fparam(self) -> None: + """Eval .pt2 without fparam should match eval with explicit default value.""" + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(self.pt2_path) + + # Eval WITHOUT fparam — model should use default (0.5) + e_no, f_no, v_no = dp.eval(_coord, _box, _atype) + # Eval WITH explicit default value + e_ex, f_ex, v_ex = dp.eval(_coord, _box, _atype, fparam=[0.5]) + + self.assertTrue(np.isfinite(e_no).all()) + self.assertTrue(np.isfinite(f_no).all()) + self.assertTrue(np.isfinite(v_no).all()) + np.testing.assert_allclose(e_no, e_ex, rtol=0, atol=1e-10) + np.testing.assert_allclose(f_no, f_ex, rtol=0, atol=1e-10) + np.testing.assert_allclose(v_no, v_ex, rtol=0, atol=1e-10) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_finetune.py b/source/tests/pt_expt/test_finetune.py index 250ba85d46..ec24e2ddd0 100644 --- a/source/tests/pt_expt/test_finetune.py +++ b/source/tests/pt_expt/test_finetune.py @@ -539,18 +539,46 @@ class TestFinetuneCLI(unittest.TestCase): @classmethod def setUpClass(cls) -> None: + from .conftest import ( + _pop_device_contexts, + ) + + _pop_device_contexts() + data_dir = os.path.join(EXAMPLE_DIR, "data") if not os.path.isdir(data_dir): raise unittest.SkipTest(f"Example data not found: {data_dir}") cls.data_dir = data_dir - def _train_pretrained(self, config: dict, tmpdir: str) -> str: - """Train a 1-step model and return checkpoint path.""" - trainer = get_trainer(config) - trainer.run() - ckpt = os.path.join(tmpdir, "model.ckpt.pt") - self.assertTrue(os.path.exists(ckpt), "Pretrained checkpoint not found") - return ckpt + # Train once and freeze to .pte and .pt2 for reuse across tests + cls.shared_tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_shared_") + old_cwd = os.getcwd() + os.chdir(cls.shared_tmpdir) + try: + config = _make_config(data_dir, model_se_e2_a, numb_steps=1) + config = update_deepmd_input(config, warning=False) + config = normalize(config) + trainer = get_trainer(config) + trainer.run() + finally: + os.chdir(old_cwd) + + cls.shared_ckpt_path = os.path.join(cls.shared_tmpdir, "model.ckpt.pt") + state = torch.load(cls.shared_ckpt_path, map_location=DEVICE, weights_only=True) + cls.shared_model_state = state["model"] if "model" in state else state + + from deepmd.pt_expt.entrypoints.main import ( + freeze, + ) + + cls.shared_pte_path = os.path.join(cls.shared_tmpdir, "shared.pte") + cls.shared_pt2_path = os.path.join(cls.shared_tmpdir, "shared.pt2") + freeze(model=cls.shared_ckpt_path, output=cls.shared_pte_path) + freeze(model=cls.shared_ckpt_path, output=cls.shared_pt2_path) + + @classmethod + def tearDownClass(cls) -> None: + shutil.rmtree(cls.shared_tmpdir, ignore_errors=True) def _assert_inherited_weights_match( self, @@ -602,15 +630,8 @@ def test_finetune_cli(self) -> None: old_cwd = os.getcwd() os.chdir(tmpdir) try: - # Phase 1: train pretrained model - config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) - config = update_deepmd_input(config, warning=False) - config = normalize(config) - ckpt_path = self._train_pretrained(config, tmpdir) - - # Save original bias - state = torch.load(ckpt_path, map_location=DEVICE, weights_only=True) - model_state = state["model"] if "model" in state else state + # Load shared pretrained model state + model_state = self.shared_model_state original_model = get_model(model_state["_extra_state"]["model_params"]).to( DEVICE ) @@ -618,7 +639,7 @@ def test_finetune_cli(self) -> None: original_wrapper.load_state_dict(model_state) original_bias = to_numpy_array(original_model.get_out_bias()).copy() - # Phase 2: finetune via CLI (lr=0 so weights stay unchanged) + # Finetune via CLI (lr=0 so weights stay unchanged) ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) ft_config["learning_rate"]["start_lr"] = 1e-30 ft_config["learning_rate"]["stop_lr"] = 1e-30 @@ -631,7 +652,7 @@ def test_finetune_cli(self) -> None: "train", ft_config_file, "--finetune", - ckpt_path, + self.shared_ckpt_path, "--skip-neighbor-stat", ] ) @@ -673,13 +694,7 @@ def test_finetune_cli_use_pretrain_script(self) -> None: old_cwd = os.getcwd() os.chdir(tmpdir) try: - # Phase 1: train pretrained model - config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) - config = update_deepmd_input(config, warning=False) - config = normalize(config) - ckpt_path = self._train_pretrained(config, tmpdir) - - # Phase 2: finetune with --use-pretrain-script + # Finetune with --use-pretrain-script # Use a config with different descriptor neuron sizes ft_model_params = deepcopy(model_se_e2_a) ft_model_params["descriptor"]["neuron"] = [4, 8] # different @@ -693,7 +708,7 @@ def test_finetune_cli_use_pretrain_script(self) -> None: "train", ft_config_file, "--finetune", - ckpt_path, + self.shared_ckpt_path, "--use-pretrain-script", "--skip-neighbor-stat", ] @@ -717,13 +732,9 @@ def test_finetune_random_fitting(self) -> None: old_cwd = os.getcwd() os.chdir(tmpdir) try: - # Phase 1: train pretrained model - config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) - config = update_deepmd_input(config, warning=False) - config = normalize(config) - ckpt_path = self._train_pretrained(config, tmpdir) + ckpt_path = self.shared_ckpt_path - # Phase 2: finetune with RANDOM (random fitting) + # Finetune with RANDOM (random fitting) ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) ft_config = update_deepmd_input(ft_config, warning=False) ft_config = normalize(ft_config) @@ -742,32 +753,18 @@ def test_finetune_random_fitting(self) -> None: finetune_links=finetune_links, ) - # Load pretrained weights for comparison - pretrained_state = torch.load( - ckpt_path, map_location=DEVICE, weights_only=True - ) - if "model" in pretrained_state: - pretrained_state = pretrained_state["model"] - pretrained_model = get_model( - pretrained_state["_extra_state"]["model_params"] - ).to(DEVICE) - pretrained_wrapper = ModelWrapper(pretrained_model) - pretrained_wrapper.load_state_dict(pretrained_state) - # Descriptor weights should match; fitting should NOT ft_state = trainer_ft.wrapper.state_dict() - pre_state = pretrained_wrapper.state_dict() self._assert_inherited_weights_match( - ft_state, pre_state, random_fitting=True + ft_state, self.shared_model_state, random_fitting=True ) finally: os.chdir(old_cwd) shutil.rmtree(tmpdir, ignore_errors=True) def test_finetune_from_pte(self) -> None: - """Train -> freeze to .pte -> finetune from .pte -> verify checkpoint.""" + """Finetune from shared .pte -> verify checkpoint.""" from deepmd.pt_expt.entrypoints.main import ( - freeze, main, ) @@ -775,18 +772,7 @@ def test_finetune_from_pte(self) -> None: old_cwd = os.getcwd() os.chdir(tmpdir) try: - # Phase 1: train pretrained model - config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) - config = update_deepmd_input(config, warning=False) - config = normalize(config) - ckpt_path = self._train_pretrained(config, tmpdir) - - # Phase 2: freeze to .pte - pte_path = os.path.join(tmpdir, "frozen.pte") - freeze(model=ckpt_path, output=pte_path) - self.assertTrue(os.path.exists(pte_path)) - - # Phase 3: finetune from .pte via CLI (lr=0 so weights stay unchanged) + # Finetune from shared .pte via CLI (lr=0 so weights stay unchanged) ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) ft_config["learning_rate"]["start_lr"] = 1e-30 ft_config["learning_rate"]["stop_lr"] = 1e-30 @@ -799,7 +785,7 @@ def test_finetune_from_pte(self) -> None: "train", ft_config_file, "--finetune", - pte_path, + self.shared_pte_path, "--skip-neighbor-stat", ] ) @@ -813,22 +799,17 @@ def test_finetune_from_pte(self) -> None: ft_model_state = ft_state["model"] if "model" in ft_state else ft_state self.assertIn("_extra_state", ft_model_state) - # Load pretrained from .pt for weight comparison - pre_state = torch.load(ckpt_path, map_location=DEVICE, weights_only=True) - pre_model_state = pre_state["model"] if "model" in pre_state else pre_state - # Inherited weights must match pretrained self._assert_inherited_weights_match( - ft_model_state, pre_model_state, random_fitting=False + ft_model_state, self.shared_model_state, random_fitting=False ) finally: os.chdir(old_cwd) shutil.rmtree(tmpdir, ignore_errors=True) def test_finetune_from_pte_use_pretrain_script(self) -> None: - """Train -> freeze to .pte -> finetune with --use-pretrain-script.""" + """Finetune from shared .pte with --use-pretrain-script.""" from deepmd.pt_expt.entrypoints.main import ( - freeze, main, ) @@ -836,17 +817,93 @@ def test_finetune_from_pte_use_pretrain_script(self) -> None: old_cwd = os.getcwd() os.chdir(tmpdir) try: - # Phase 1: train pretrained model - config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) - config = update_deepmd_input(config, warning=False) - config = normalize(config) - ckpt_path = self._train_pretrained(config, tmpdir) + # Finetune from shared .pte with --use-pretrain-script + ft_model_params = deepcopy(model_se_e2_a) + ft_model_params["descriptor"]["neuron"] = [4, 8] # different + ft_config = _make_config(self.data_dir, ft_model_params, numb_steps=1) + ft_config_file = os.path.join(tmpdir, "finetune_input.json") + with open(ft_config_file, "w") as f: + json.dump(ft_config, f) + + main( + [ + "train", + ft_config_file, + "--finetune", + self.shared_pte_path, + "--use-pretrain-script", + "--skip-neighbor-stat", + ] + ) + + # Verify the output config was updated from pretrained + with open(os.path.join(tmpdir, "out.json")) as f: + output_config = json.load(f) + # Descriptor neuron should be from pretrained, not from ft_config + self.assertEqual( + output_config["model"]["descriptor"]["neuron"], + model_se_e2_a["descriptor"]["neuron"], + ) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_finetune_from_pt2(self) -> None: + """Finetune from shared .pt2 -> verify checkpoint.""" + from deepmd.pt_expt.entrypoints.main import ( + main, + ) - # Phase 2: freeze to .pte (embeds model_params) - pte_path = os.path.join(tmpdir, "frozen.pte") - freeze(model=ckpt_path, output=pte_path) + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pt2_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + # Finetune from shared .pt2 via CLI (lr=0 so weights stay unchanged) + ft_config = _make_config(self.data_dir, model_se_e2_a, numb_steps=1) + ft_config["learning_rate"]["start_lr"] = 1e-30 + ft_config["learning_rate"]["stop_lr"] = 1e-30 + ft_config_file = os.path.join(tmpdir, "finetune_input.json") + with open(ft_config_file, "w") as f: + json.dump(ft_config, f) - # Phase 3: finetune from .pte with --use-pretrain-script + main( + [ + "train", + ft_config_file, + "--finetune", + self.shared_pt2_path, + "--skip-neighbor-stat", + ] + ) + + # Verify new checkpoint exists + ft_ckpt = os.path.join(tmpdir, "model.ckpt.pt") + self.assertTrue(os.path.exists(ft_ckpt), "Finetune checkpoint not found") + + # Load finetuned model and verify it's valid + ft_state = torch.load(ft_ckpt, map_location=DEVICE, weights_only=True) + ft_model_state = ft_state["model"] if "model" in ft_state else ft_state + self.assertIn("_extra_state", ft_model_state) + + # Inherited weights must match pretrained + self._assert_inherited_weights_match( + ft_model_state, self.shared_model_state, random_fitting=False + ) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_finetune_from_pt2_use_pretrain_script(self) -> None: + """Finetune from shared .pt2 with --use-pretrain-script.""" + from deepmd.pt_expt.entrypoints.main import ( + main, + ) + + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ft_pt2_ups_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + # Finetune from shared .pt2 with --use-pretrain-script ft_model_params = deepcopy(model_se_e2_a) ft_model_params["descriptor"]["neuron"] = [4, 8] # different ft_config = _make_config(self.data_dir, ft_model_params, numb_steps=1) @@ -859,7 +916,7 @@ def test_finetune_from_pte_use_pretrain_script(self) -> None: "train", ft_config_file, "--finetune", - pte_path, + self.shared_pt2_path, "--use-pretrain-script", "--skip-neighbor-stat", ]