Merge branch 'devel' into D0516_dynamic_sel

Author: iProzd

deepmd/pt/model/network/network.py

@@ -286,7 +286,7 @@ def forward(self, atype):
         type_embedding:
 
         """
-        return self.embedding(atype.device)[atype]
+        return torch.embedding(self.embedding(atype.device), atype)
 
     def get_full_embedding(self, device: torch.device):
         """

deepmd/pt/train/training.py

@@ -2,6 +2,9 @@
 import functools
 import logging
 import time
+from collections.abc import (
+    Iterable,
+)
 from copy import (
     deepcopy,
 )
@@ -47,7 +50,6 @@
     dp_random,
 )
 from deepmd.pt.utils.dataloader import (
-    BufferedIterator,
     get_sampler_from_params,
 )
 from deepmd.pt.utils.env import (
@@ -159,8 +161,24 @@ def get_opt_param(params):
             }
             return opt_type, opt_param
 
+        def cycle_iterator(iterable: Iterable):
+            """
+            Produces an infinite iterator by repeatedly cycling through the given iterable.
+
+            Args:
+                iterable (Iterable): The iterable to cycle through.
+
+            Yields
+            ------
+            Any: The next item from the iterable, cycling back to the beginning when the end is reached.
+            """
+            while True:
+                with torch.device("cpu"):
+                    it = iter(iterable)
+                yield from it
+
         def get_data_loader(_training_data, _validation_data, _training_params):
-            def get_dataloader_and_buffer(_data, _params):
+            def get_dataloader_and_iter(_data, _params):
                 _sampler = get_sampler_from_params(_data, _params)
                 if _sampler is None:
                     log.warning(
@@ -177,33 +195,32 @@ def get_dataloader_and_buffer(_data, _params):
                     collate_fn=lambda batch: batch,  # prevent extra conversion
                     pin_memory=True,
                 )
-                with torch.device("cpu"):
-                    _data_buffered = BufferedIterator(iter(_dataloader))
-                return _dataloader, _data_buffered
+                _data_iter = cycle_iterator(_dataloader)
+                return _dataloader, _data_iter
 
-            training_dataloader, training_data_buffered = get_dataloader_and_buffer(
+            training_dataloader, training_data_iter = get_dataloader_and_iter(
                 _training_data, _training_params["training_data"]
             )
 
             if _validation_data is not None:
                 (
                     validation_dataloader,
-                    validation_data_buffered,
-                ) = get_dataloader_and_buffer(
+                    validation_data_iter,
+                ) = get_dataloader_and_iter(
                     _validation_data, _training_params["validation_data"]
                 )
                 valid_numb_batch = _training_params["validation_data"].get(
                     "numb_btch", 1
                 )
             else:
                 validation_dataloader = None
-                validation_data_buffered = None
+                validation_data_iter = None
                 valid_numb_batch = 1
             return (
                 training_dataloader,
-                training_data_buffered,
+                training_data_iter,
                 validation_dataloader,
-                validation_data_buffered,
+                validation_data_iter,
                 valid_numb_batch,
             )
 
@@ -1064,48 +1081,15 @@ def save_model(self, save_path, lr=0.0, step=0) -> None:
             checkpoint_files[0].unlink()
 
     def get_data(self, is_train=True, task_key="Default"):
-        if not self.multi_task:
-            if is_train:
-                try:
-                    batch_data = next(iter(self.training_data))
-                except StopIteration:
-                    # Refresh the status of the dataloader to start from a new epoch
-                    with torch.device("cpu"):
-                        self.training_data = BufferedIterator(
-                            iter(self.training_dataloader)
-                        )
-                    batch_data = next(iter(self.training_data))
-            else:
-                if self.validation_data is None:
-                    return {}, {}, {}
-                try:
-                    batch_data = next(iter(self.validation_data))
-                except StopIteration:
-                    self.validation_data = BufferedIterator(
-                        iter(self.validation_dataloader)
-                    )
-                    batch_data = next(iter(self.validation_data))
+        if is_train:
+            iterator = self.training_data
         else:
-            if is_train:
-                try:
-                    batch_data = next(iter(self.training_data[task_key]))
-                except StopIteration:
-                    # Refresh the status of the dataloader to start from a new epoch
-                    self.training_data[task_key] = BufferedIterator(
-                        iter(self.training_dataloader[task_key])
-                    )
-                    batch_data = next(iter(self.training_data[task_key]))
-            else:
-                if self.validation_data[task_key] is None:
-                    return {}, {}, {}
-                try:
-                    batch_data = next(iter(self.validation_data[task_key]))
-                except StopIteration:
-                    self.validation_data[task_key] = BufferedIterator(
-                        iter(self.validation_dataloader[task_key])
-                    )
-                    batch_data = next(iter(self.validation_data[task_key]))
-
+            iterator = self.validation_data
+        if self.multi_task:
+            iterator = iterator[task_key]
+        if iterator is None:
+            return {}, {}, {}
+        batch_data = next(iterator)
         for key in batch_data.keys():
             if key == "sid" or key == "fid" or key == "box" or "find_" in key:
                 continue

deepmd/pt/utils/dataloader.py

@@ -1,16 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
 import os
-import time
 from multiprocessing.dummy import (
     Pool,
 )
-from queue import (
-    Queue,
-)
-from threading import (
-    Thread,
-)
 
 import h5py
 import numpy as np
@@ -173,7 +166,9 @@ def construct_dataset(system):
                 num_workers=0,  # Should be 0 to avoid too many threads forked
                 sampler=system_sampler,
                 collate_fn=collate_batch,
-                shuffle=(not (dist.is_available() and dist.is_initialized()))
+                shuffle=(
+                    not (dist.is_available() and dist.is_initialized())
+                )  # distributed sampler will do the shuffling by default
                 and shuffle,
             )
             self.dataloaders.append(system_dataloader)
@@ -200,11 +195,12 @@ def __len__(self) -> int:
 
     def __getitem__(self, idx):
         # log.warning(str(torch.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx]))
-        try:
-            batch = next(self.iters[idx])
-        except StopIteration:
-            self.iters[idx] = iter(self.dataloaders[idx])
-            batch = next(self.iters[idx])
+        with torch.device("cpu"):
+            try:
+                batch = next(self.iters[idx])
+            except StopIteration:
+                self.iters[idx] = iter(self.dataloaders[idx])
+                batch = next(self.iters[idx])
         batch["sid"] = idx
         return batch
 
@@ -235,54 +231,6 @@ def print_summary(
             )
 
 
-class BackgroundConsumer(Thread):
-    def __init__(self, queue, source) -> None:
-        super().__init__()
-        self.daemon = True
-        self._queue = queue
-        self._source = source  # Main DL iterator
-
-    def run(self) -> None:
-        for item in self._source:
-            self._queue.put(item)  # Blocking if the queue is full
-
-        # Signal the consumer we are done; this should not happen for DataLoader
-        self._queue.put(StopIteration())
-
-
-QUEUESIZE = 32
-
-
-class BufferedIterator:
-    def __init__(self, iterable) -> None:
-        self._queue = Queue(QUEUESIZE)
-        self._iterable = iterable
-        self._consumer = BackgroundConsumer(self._queue, self._iterable)
-        self._consumer.start()
-        self.last_warning_time = time.time()
-
-    def __iter__(self):
-        return self
-
-    def __len__(self) -> int:
-        return len(self._iterable)
-
-    def __next__(self):
-        start_wait = time.time()
-        item = self._queue.get()
-        wait_time = time.time() - start_wait
-        if (
-            wait_time > 1.0 and start_wait - self.last_warning_time > 15 * 60
-        ):  # Even for Multi-Task training, each step usually takes < 1s
-            log.warning(
-                f"Data loading is slow, waited {wait_time:.2f} seconds. Ignoring this warning for 15 minutes."
-            )
-            self.last_warning_time = start_wait
-        if isinstance(item, Exception):
-            raise item
-        return item
-
-
 def collate_batch(batch):
     example = batch[0]
     result = {}
@@ -320,7 +268,11 @@ def get_weighted_sampler(training_data, prob_style, sys_prob=False):
     # training_data.total_batch is the size of one epoch, you can increase it to avoid too many  rebuilding of iterators
     len_sampler = training_data.total_batch * max(env.NUM_WORKERS, 1)
     with torch.device("cpu"):
-        sampler = WeightedRandomSampler(probs, len_sampler, replacement=True)
+        sampler = WeightedRandomSampler(
+            probs,
+            len_sampler,
+            replacement=True,
+        )
     return sampler
 
 

source/tests/pt/model/test_saveload_dpa1.py

@@ -25,7 +25,6 @@
     env,
 )
 from deepmd.pt.utils.dataloader import (
-    BufferedIterator,
     DpLoaderSet,
 )
 from deepmd.pt.utils.stat import (
@@ -72,8 +71,13 @@ def setUp(self) -> None:
             drop_last=False,
             pin_memory=True,
         )
+
+        def cycle_iterator(iterable):
+            while True:
+                yield from iterable
+
         with torch.device("cpu"):
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
+            self.training_data = cycle_iterator(self.training_dataloader)
         self.loss = EnergyStdLoss(**self.config["loss"])
         self.cur_lr = 1
         self.task_key = "Default"
@@ -111,12 +115,8 @@ def create_wrapper(self, read: bool):
         return ModelWrapper(model, self.loss)
 
     def get_data(self):
-        try:
-            batch_data = next(iter(self.training_data))
-        except StopIteration:
-            # Refresh the status of the dataloader to start from a new epoch
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
-            batch_data = next(iter(self.training_data))
+        with torch.device("cpu"):
+            batch_data = next(self.training_data)
         input_dict = {}
         for item in ["coord", "atype", "box"]:
             if item in batch_data:

source/tests/pt/model/test_saveload_se_e2_a.py

@@ -25,7 +25,6 @@
     env,
 )
 from deepmd.pt.utils.dataloader import (
-    BufferedIterator,
     DpLoaderSet,
 )
 from deepmd.pt.utils.stat import (
@@ -72,8 +71,13 @@ def setUp(self) -> None:
             drop_last=False,
             pin_memory=True,
         )
+
+        def cycle_iterator(iterable):
+            while True:
+                yield from iterable
+
         with torch.device("cpu"):
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
+            self.training_data = cycle_iterator(self.training_dataloader)
         self.loss = EnergyStdLoss(**self.config["loss"])
         self.cur_lr = 1
         self.task_key = "Default"
@@ -105,12 +109,8 @@ def create_wrapper(self):
         return ModelWrapper(model, self.loss)
 
     def get_data(self):
-        try:
-            batch_data = next(iter(self.training_data))
-        except StopIteration:
-            # Refresh the status of the dataloader to start from a new epoch
-            self.training_data = BufferedIterator(iter(self.training_dataloader))
-            batch_data = next(iter(self.training_data))
+        with torch.device("cpu"):
+            batch_data = next(self.training_data)
         input_dict = {}
         for item in ["coord", "atype", "box"]:
             if item in batch_data: