From c267bd3a0c4ac67bf1631676d1c3e7a5bfb38de3 Mon Sep 17 00:00:00 2001 From: ExtReMLapin <3909752+ExtReMLapin@users.noreply.github.com> Date: Fri, 19 Jan 2024 16:49:54 +0100 Subject: [PATCH] Fixed dataloader CPU bottleneck for small batch sizes (#7659) Co-authored-by: Glenn Jocher --- ultralytics/data/build.py | 2 +- ultralytics/engine/trainer.py | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/ultralytics/data/build.py b/ultralytics/data/build.py index 4c6a96d1..f5c4219f 100644 --- a/ultralytics/data/build.py +++ b/ultralytics/data/build.py @@ -107,7 +107,7 @@ def build_dataloader(dataset, batch, workers, shuffle=True, rank=-1): """Return an InfiniteDataLoader or DataLoader for training or validation set.""" batch = min(batch, len(dataset)) nd = torch.cuda.device_count() # number of CUDA devices - nw = min([os.cpu_count() // max(nd, 1), batch, workers]) # number of workers + nw = min([os.cpu_count() // max(nd, 1), workers]) # number of workers sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle) generator = torch.Generator() generator.manual_seed(6148914691236517205 + RANK) diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py index 0fa98396..75b3ea3f 100644 --- a/ultralytics/engine/trainer.py +++ b/ultralytics/engine/trainer.py @@ -332,10 +332,7 @@ class BaseTrainer: f'Image sizes {self.args.imgsz} train, {self.args.imgsz} val\n' f'Using {self.train_loader.num_workers * (world_size or 1)} dataloader workers\n' f"Logging results to {colorstr('bold', self.save_dir)}\n" - f'Starting training for ' - f'{self.args.time} hours...' - if self.args.time - else f"{self.epochs} epochs..." + f'Starting training for ' + (f"{self.args.time} hours..." if self.args.time else f"{self.epochs} epochs...") ) if self.args.close_mosaic: base_idx = (self.epochs - self.args.close_mosaic) * nb