Skip to content

Commit

Permalink
revert changes on using process for no significant improvement
Browse files Browse the repository at this point in the history
  • Loading branch information
caic99 committed Dec 2, 2024
1 parent 68dd750 commit 5a454d5
Showing 1 changed file with 9 additions and 15 deletions.
24 changes: 9 additions & 15 deletions deepmd/pt/utils/dataloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@
import logging
import os
import time
from functools import (
partial,
)
from multiprocessing import (

from multiprocessing.dummy import (
Pool,
)
from queue import (
Expand Down Expand Up @@ -57,13 +55,6 @@ def setup_seed(seed) -> None:
torch.backends.cudnn.deterministic = True


def construct_dataset(system, type_map):
return DeepmdDataSetForLoader(
system=system,
type_map=type_map,
)


class DpLoaderSet(Dataset):
"""A dataset for storing DataLoaders to multiple Systems.
Expand Down Expand Up @@ -99,7 +90,11 @@ def __init__(
if len(systems) >= 100:
log.info(f"Constructing DataLoaders from {len(systems)} systems")

construct_dataset_systems = partial(construct_dataset, type_map=type_map)
def construct_dataset(system):
return DeepmdDataSetForLoader(
system=system,
type_map=type_map,
)

with Pool(
os.cpu_count()
Expand All @@ -109,7 +104,7 @@ def __init__(
else 1
)
) as pool:
self.systems = pool.map(construct_dataset_systems, systems)
self.systems = pool.map(construct_dataset, systems)

self.sampler_list: list[DistributedSampler] = []
self.index = []
Expand Down Expand Up @@ -235,13 +230,12 @@ def __init__(self, iterable) -> None:
self._consumer = BackgroundConsumer(self._queue, self._iterable)
self._consumer.start()
self.last_warning_time = time.time()
self.len = len(iterable)

def __iter__(self):
return self

def __len__(self) -> int:
return self.len
return len(self._iterable)

def __next__(self):
start_wait = time.time()
Expand Down

0 comments on commit 5a454d5

Please sign in to comment.