From 337ca3490fa7879ceeeadf6c2b73d67504ff4b4f Mon Sep 17 00:00:00 2001 From: Gabriele <34710621+Gabrysse@users.noreply.github.com> Date: Wed, 1 Nov 2023 08:55:51 -0700 Subject: [PATCH] Bug fix with prefetch_factor Summary: Fix https://github.com/facebookresearch/detectron2/issues/5086. The PyTorch DataLoader class can accepts "None" as prefetch_factor in recent versions (> 2.0). For backward compatibility, it is better to set a default value, specifically 2 as in previous PyTorch versions. Looking at the more recent DataLoader source code, it sets the value 2 if prefetch_factor is found to be None. Pull Request resolved: https://github.com/facebookresearch/detectron2/pull/5091 Reviewed By: wat3rBro Differential Revision: D50693761 Pulled By: ezyang fbshipit-source-id: 479ec794009be9e95d27c401143a88dcd45a6eff --- detectron2/data/build.py | 20 +++++--------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/detectron2/data/build.py b/detectron2/data/build.py index 5d23f616b4..87b4a0050b 100644 --- a/detectron2/data/build.py +++ b/detectron2/data/build.py @@ -300,9 +300,7 @@ def build_batch_data_loader( num_workers=0, collate_fn=None, drop_last: bool = True, - prefetch_factor=None, - persistent_workers=False, - pin_memory=False, + **kwargs, ): """ Build a batched dataloader. The main differences from `torch.utils.data.DataLoader` are: @@ -341,9 +339,7 @@ def build_batch_data_loader( num_workers=num_workers, collate_fn=operator.itemgetter(0), # don't batch, but yield individual elements worker_init_fn=worker_init_reset_seed, - prefetch_factor=prefetch_factor, - persistent_workers=persistent_workers, - pin_memory=pin_memory, + **kwargs ) # yield individual mapped dict data_loader = AspectRatioGroupedDataset(data_loader, batch_size) if collate_fn is None: @@ -357,9 +353,7 @@ def build_batch_data_loader( num_workers=num_workers, collate_fn=trivial_batch_collator if collate_fn is None else collate_fn, worker_init_fn=worker_init_reset_seed, - prefetch_factor=prefetch_factor, - persistent_workers=persistent_workers, - pin_memory=pin_memory, + **kwargs ) @@ -499,9 +493,7 @@ def build_detection_train_loader( aspect_ratio_grouping=True, num_workers=0, collate_fn=None, - prefetch_factor=None, - persistent_workers=False, - pin_memory=False, + **kwargs ): """ Build a dataloader for object detection with some default features. @@ -553,9 +545,7 @@ def build_detection_train_loader( aspect_ratio_grouping=aspect_ratio_grouping, num_workers=num_workers, collate_fn=collate_fn, - prefetch_factor=prefetch_factor, - persistent_workers=persistent_workers, - pin_memory=pin_memory, + **kwargs )