diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index e30c97bd98..a5817727f5 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -90,6 +90,5 @@ def dynamic_metadata( ], "torch": [ "torch>=2a", - "tqdm", ], } diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 8ea69c8489..e4c672765b 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -14,12 +14,6 @@ import numpy as np import torch -from tqdm import ( - tqdm, -) -from tqdm.contrib.logging import ( - logging_redirect_tqdm, -) from deepmd.common import ( symlink_prefix_files, @@ -47,7 +41,6 @@ ) from deepmd.pt.utils.env import ( DEVICE, - DISABLE_TQDM, JIT, LOCAL_RANK, NUM_WORKERS, @@ -662,29 +655,24 @@ def log_loss_valid(_task_key="Default"): f.write(str(self.latest_model)) self.t0 = time.time() - with logging_redirect_tqdm(): - for step_id in tqdm( - range(self.num_steps), - disable=(bool(dist.get_rank()) if dist.is_initialized() else False) - or DISABLE_TQDM, - ): # set to None to disable on non-TTY; disable on not rank 0 - if step_id < self.start_step: - continue - if self.multi_task: - chosen_index_list = dp_random.choice( - np.arange(self.num_model), - p=np.array(self.model_prob), - size=self.world_size, - replace=True, - ) - assert chosen_index_list.size == self.world_size - model_index = chosen_index_list[self.rank] - model_key = self.model_keys[model_index] - else: - model_key = "Default" - step(step_id, model_key) - if JIT: - break + for step_id in range(self.num_steps): + if step_id < self.start_step: + continue + if self.multi_task: + chosen_index_list = dp_random.choice( + np.arange(self.num_model), + p=np.array(self.model_prob), + size=self.world_size, + replace=True, + ) + assert chosen_index_list.size == self.world_size + model_index = chosen_index_list[self.rank] + model_key = self.model_keys[model_index] + else: + model_key = "Default" + step(step_id, model_key) + if JIT: + break if ( self.rank == 0 or dist.get_rank() == 0 diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index 3920499d3a..c104e64491 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -13,9 +13,6 @@ from torch.utils.data import ( Dataset, ) -from tqdm import ( - trange, -) from deepmd.pt.utils import ( dp_random, @@ -506,7 +503,7 @@ def preprocess(self, batch): assert batch["atype"].max() < len(self._type_map) nlist, nlist_loc, nlist_type, shift, mapping = [], [], [], [], [] - for sid in trange(n_frames, disable=env.DISABLE_TQDM): + for sid in range(n_frames): region = Region3D(box[sid]) nloc = atype[sid].shape[0] _coord = normalize_coord(coord[sid], region, nloc) diff --git a/deepmd/pt/utils/env.py b/deepmd/pt/utils/env.py index 5b6eaf7c14..6fa72943c7 100644 --- a/deepmd/pt/utils/env.py +++ b/deepmd/pt/utils/env.py @@ -8,7 +8,6 @@ GLOBAL_NP_FLOAT_PRECISION = getattr(np, PRECISION) GLOBAL_PT_FLOAT_PRECISION = getattr(torch, PRECISION) GLOBAL_ENER_FLOAT_PRECISION = getattr(np, PRECISION) -DISABLE_TQDM = os.environ.get("DISABLE_TQDM", False) SAMPLER_RECORD = os.environ.get("SAMPLER_RECORD", False) try: # only linux diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index 837a0104f9..18ee4d9abe 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -3,9 +3,6 @@ import numpy as np import torch -from tqdm import ( - trange, -) from deepmd.pt.utils import ( env, @@ -40,7 +37,7 @@ def make_stat_input(datasets, dataloaders, nbatches): if datasets[0].mixed_type: keys.append("real_natoms_vec") logging.info(f"Packing data for statistics from {len(datasets)} systems") - for i in trange(len(datasets), disable=env.DISABLE_TQDM): + for i in range(len(datasets)): sys_stat = {key: [] for key in keys} iterator = iter(dataloaders[i]) for _ in range(nbatches):