adopt

Author: OutisLi

deepmd/pt/train/training.py

@@ -278,13 +278,17 @@ def resolve_model_prob(
             model_training_data: dict[str, DpLoaderSet],
         ) -> np.ndarray:
             model_prob = np.zeros(len(model_keys), dtype=np.float64)
-            if model_prob_config is not None:
+            if model_prob_config:
                 for ii, model_key in enumerate(model_keys):
                     if model_key in model_prob_config:
                         model_prob[ii] = float(model_prob_config[model_key])
             else:
                 for ii, model_key in enumerate(model_keys):
                     model_prob[ii] = float(len(model_training_data[model_key]))
+            if not np.all(np.isfinite(model_prob)):
+                raise ValueError("Model prob must be finite.")
+            if np.any(model_prob < 0.0):
+                raise ValueError("Model prob must be non-negative.")
             sum_prob = float(np.sum(model_prob))
             if sum_prob <= 0.0:
                 raise ValueError("Sum of model prob must be larger than 0!")

deepmd/utils/argcheck.py

@@ -3213,15 +3213,22 @@ def mixed_precision_args() -> list[Argument]:  # ! added by Denghui.
 def training_args(
     multi_task: bool = False,
 ) -> list[Argument]:  # ! modified by Ziyao: data configuration isolated.
-    doc_numb_steps = "Number of training batches. Each training uses one batch of data. If set, this value takes precedence over num_epoch."
+    doc_numb_steps = (
+        "Number of training batches. Each training uses one batch of data. "
+        "If set, this value takes precedence over num_epoch. If both numb_steps "
+        "and num_epoch are not set, a ValueError is raised."
+    )
     doc_num_epoch = (
-        "Number of training epochs. "
+        "Number of training epochs (can be fractional). "
         "When numb_steps is not set, the total steps are computed as "
-        "ceil(num_epoch * total_numb_batch). For each training dataset, "
-        "total_numb_batch is computed as ceil(max_i(n_bch_i / p_i)), where p_i "
-        "is the sampling probability of system i after sys_probs/auto_prob. "
-        "In multi-task mode, total_numb_batch is the model_prob-weighted sum "
-        "over tasks."
+        "ceil(num_epoch * total_numb_batch). For each task, total_numb_batch "
+        "is computed as ceil(max_i(n_bch_i / p_i)), where n_bch_i is the number "
+        "of batches for system i and p_i is the sampling probability after "
+        "sys_probs/auto_prob normalization. In multi-task mode, model_prob is "
+        "normalized to sum to 1, per-task total_numb_batch values are computed "
+        "as above, and the final total_numb_batch is their model_prob-weighted "
+        "sum. At least one of numb_steps or num_epoch must be set; otherwise a "
+        "ValueError is raised."
     )
     doc_seed = "The random seed for getting frames from the training data set."
     doc_disp_file = "The file for printing learning curve."
@@ -3295,7 +3302,11 @@ def training_args(
     args += [
         mixed_precision_data,
         Argument(
-            "numb_steps", int, optional=True, doc=doc_numb_steps, alias=["stop_batch"]
+            "numb_steps",
+            int,
+            optional=True,
+            doc=doc_numb_steps,
+            alias=["stop_batch", "num_steps"],
         ),
         Argument(
             "num_epoch",

source/tests/pt/test_sampler.py

@@ -17,11 +17,6 @@
 from deepmd.pt.utils import (
     dp_random,
 )
-from deepmd.pt.utils.dataloader import (
-    DpLoaderSet,
-    get_sampler_from_params,
-    get_weighted_sampler,
-)
 from deepmd.tf.common import (
     expand_sys_str,
 )
@@ -67,7 +62,7 @@ def setUp(self) -> None:
         self.systems = config["training"]["validation_data"]["systems"]
         if isinstance(self.systems, str):
             self.systems = expand_sys_str(self.systems)
-        self.my_dataset = DpLoaderSet(
+        self.my_dataset = pt_dataloader.DpLoaderSet(
             self.systems,
             self.batch_size,
             self.type_map,
@@ -81,7 +76,9 @@ def setUp(self) -> None:
     def tearDown(self) -> None:
         self._monkeypatch.undo()
 
-    def _make_dataloader(self, dataset: DpLoaderSet, sampler) -> DataLoader:
+    def _make_dataloader(
+        self, dataset: pt_dataloader.DpLoaderSet, sampler
+    ) -> DataLoader:
         return DataLoader(
             dataset,
             sampler=sampler,
@@ -96,6 +93,18 @@ def _normalize_probs(self, weights: np.ndarray) -> np.ndarray:
         return weights / np.sum(weights)
 
     def _compute_total_numb_batch(self, nbatches: np.ndarray, probs: np.ndarray) -> int:
+        # NOTE: This is a simplified test-only variant of training.py logic.
+        nbatches = np.asarray(nbatches, dtype=np.float64)
+        probs = np.asarray(probs, dtype=np.float64)
+        if nbatches.shape != probs.shape:
+            raise ValueError(
+                "nbatches and probs must have the same shape in this test helper."
+            )
+        if not np.all(probs > 0.0):
+            raise ValueError(
+                "Zero or negative sampling probabilities are not supported in this "
+                "test helper."
+            )
         return int(np.ceil(np.max(nbatches / probs)))
 
     def _sample_sid_counts(
@@ -156,7 +165,9 @@ def _sample_multitask_counts(
     def test_sampler_debug_info(self) -> None:
         dataloader = DataLoader(
             self.my_dataset,
-            sampler=get_weighted_sampler(self.my_dataset, prob_style="prob_sys_size"),
+            sampler=pt_dataloader.get_weighted_sampler(
+                self.my_dataset, prob_style="prob_sys_size"
+            ),
             batch_size=None,
             num_workers=0,  # setting to 0 diverges the behavior of its iterator; should be >=1
             drop_last=False,
@@ -171,31 +182,37 @@ def test_sampler_debug_info(self) -> None:
 
     def test_auto_prob_uniform(self) -> None:
         auto_prob_style = "prob_uniform"
-        sampler = get_weighted_sampler(self.my_dataset, prob_style=auto_prob_style)
+        sampler = pt_dataloader.get_weighted_sampler(
+            self.my_dataset, prob_style=auto_prob_style
+        )
         my_probs = np.array(sampler.weights)
         self.dp_dataset.set_sys_probs(auto_prob_style=auto_prob_style)
         dp_probs = np.array(self.dp_dataset.sys_probs)
         self.assertTrue(np.allclose(my_probs, dp_probs))
 
     def test_auto_prob_sys_size(self) -> None:
         auto_prob_style = "prob_sys_size"
-        sampler = get_weighted_sampler(self.my_dataset, prob_style=auto_prob_style)
+        sampler = pt_dataloader.get_weighted_sampler(
+            self.my_dataset, prob_style=auto_prob_style
+        )
         my_probs = np.array(sampler.weights)
         self.dp_dataset.set_sys_probs(auto_prob_style=auto_prob_style)
         dp_probs = np.array(self.dp_dataset.sys_probs)
         self.assertTrue(np.allclose(my_probs, dp_probs))
 
     def test_auto_prob_sys_size_ext(self) -> None:
         auto_prob_style = "prob_sys_size;0:1:0.2;1:3:0.8"
-        sampler = get_weighted_sampler(self.my_dataset, prob_style=auto_prob_style)
+        sampler = pt_dataloader.get_weighted_sampler(
+            self.my_dataset, prob_style=auto_prob_style
+        )
         my_probs = np.array(sampler.weights)
         self.dp_dataset.set_sys_probs(auto_prob_style=auto_prob_style)
         dp_probs = np.array(self.dp_dataset.sys_probs)
         self.assertTrue(np.allclose(my_probs, dp_probs))
 
     def test_sys_probs(self) -> None:
         sys_probs = [0.1, 0.4, 0.5]
-        sampler = get_weighted_sampler(
+        sampler = pt_dataloader.get_weighted_sampler(
             self.my_dataset, prob_style=sys_probs, sys_prob=True
         )
         my_probs = np.array(sampler.weights)
@@ -209,7 +226,7 @@ def test_sys_probs_end2end(self):
             "sys_probs": sys_probs,
             "auto_prob": "prob_sys_size",
         }  # use sys_probs first
-        sampler = get_sampler_from_params(self.my_dataset, _params)
+        sampler = pt_dataloader.get_sampler_from_params(self.my_dataset, _params)
         my_probs = np.array(sampler.weights)
         self.dp_dataset.set_sys_probs(sys_probs=sys_probs)
         dp_probs = np.array(self.dp_dataset.sys_probs)
@@ -218,7 +235,7 @@ def test_sys_probs_end2end(self):
     def test_auto_prob_sys_size_ext_end2end(self):
         auto_prob_style = "prob_sys_size;0:1:0.2;1:3:0.8"
         _params = {"sys_probs": None, "auto_prob": auto_prob_style}  # use auto_prob
-        sampler = get_sampler_from_params(self.my_dataset, _params)
+        sampler = pt_dataloader.get_sampler_from_params(self.my_dataset, _params)
         my_probs = np.array(sampler.weights)
         self.dp_dataset.set_sys_probs(auto_prob_style=auto_prob_style)
         dp_probs = np.array(self.dp_dataset.sys_probs)
@@ -231,7 +248,7 @@ def test_sampling_stability_single_task(self) -> None:
             str(Path(__file__).parent / "water/data/data_1"),
             str(Path(__file__).parent / "water/data/single"),
         ]
-        dataset_epoch = DpLoaderSet(
+        dataset_epoch = pt_dataloader.DpLoaderSet(
             systems,
             self.batch_size,
             self.type_map,
@@ -240,7 +257,7 @@ def test_sampling_stability_single_task(self) -> None:
         )
         sys_probs = [0.2, 0.3, 0.5]
         params = {"sys_probs": sys_probs, "auto_prob": "prob_sys_size"}
-        sampler_epoch = get_sampler_from_params(dataset_epoch, params)
+        sampler_epoch = pt_dataloader.get_sampler_from_params(dataset_epoch, params)
         probs = self._normalize_probs(np.asarray(sampler_epoch.weights))
         nbatches = np.asarray(dataset_epoch.index, dtype=np.float64)
         total_numb_batch = self._compute_total_numb_batch(nbatches, probs)
@@ -257,14 +274,14 @@ def test_sampling_stability_single_task(self) -> None:
         self.assertTrue(np.allclose(empirical_epoch, probs, atol=0.1))
 
         # === Step 3. Sample Using Explicit Steps ===
-        dataset_steps = DpLoaderSet(
+        dataset_steps = pt_dataloader.DpLoaderSet(
             systems,
             self.batch_size,
             self.type_map,
             seed=10,
             shuffle=False,
         )
-        sampler_steps = get_sampler_from_params(dataset_steps, params)
+        sampler_steps = pt_dataloader.get_sampler_from_params(dataset_steps, params)
         torch.manual_seed(123)
         dataloader_steps = self._make_dataloader(dataset_steps, sampler_steps)
         counts_steps = self._sample_sid_counts(
@@ -283,24 +300,24 @@ def test_sampling_stability_multi_task(self) -> None:
             str(Path(__file__).parent / "water/data/data_1"),
             str(Path(__file__).parent / "water/data/single"),
         ]
-        dataset_1 = DpLoaderSet(
+        dataset_1 = pt_dataloader.DpLoaderSet(
             systems_1,
             self.batch_size,
             self.type_map,
             seed=10,
             shuffle=False,
         )
-        dataset_2 = DpLoaderSet(
+        dataset_2 = pt_dataloader.DpLoaderSet(
             systems_2,
             self.batch_size,
             self.type_map,
             seed=10,
             shuffle=False,
         )
-        sampler_1 = get_sampler_from_params(
+        sampler_1 = pt_dataloader.get_sampler_from_params(
             dataset_1, {"sys_probs": [0.7, 0.3], "auto_prob": "prob_sys_size"}
         )
-        sampler_2 = get_sampler_from_params(
+        sampler_2 = pt_dataloader.get_sampler_from_params(
             dataset_2, {"sys_probs": [0.4, 0.6], "auto_prob": "prob_sys_size"}
         )
         probs_1 = self._normalize_probs(np.asarray(sampler_1.weights))
@@ -352,24 +369,24 @@ def test_sampling_stability_multi_task(self) -> None:
         )
 
         # === Step 3. Sample Using Explicit Steps ===
-        dataset_1b = DpLoaderSet(
+        dataset_1b = pt_dataloader.DpLoaderSet(
             systems_1,
             self.batch_size,
             self.type_map,
             seed=10,
             shuffle=False,
         )
-        dataset_2b = DpLoaderSet(
+        dataset_2b = pt_dataloader.DpLoaderSet(
             systems_2,
             self.batch_size,
             self.type_map,
             seed=10,
             shuffle=False,
         )
-        sampler_1b = get_sampler_from_params(
+        sampler_1b = pt_dataloader.get_sampler_from_params(
             dataset_1b, {"sys_probs": [0.7, 0.3], "auto_prob": "prob_sys_size"}
         )
-        sampler_2b = get_sampler_from_params(
+        sampler_2b = pt_dataloader.get_sampler_from_params(
             dataset_2b, {"sys_probs": [0.4, 0.6], "auto_prob": "prob_sys_size"}
         )
         dataloaders_steps = {