From c91d1db7161f4cffe70535378b81faf3ff4549b4 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 1 Sep 2022 14:30:21 +0200 Subject: [PATCH] Update dataloaders.py (#9250) Signed-off-by: Glenn Jocher Signed-off-by: Glenn Jocher --- utils/dataloaders.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/utils/dataloaders.py b/utils/dataloaders.py index 84215925284e..a4e6c0cfef18 100755 --- a/utils/dataloaders.py +++ b/utils/dataloaders.py @@ -40,6 +40,7 @@ VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' # include video suffixes BAR_FORMAT = '{l_bar}{bar:10}{r_bar}{bar:-10b}' # tqdm bar format LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html +PIN_MEMORY = str(os.getenv('PIN_MEMORY', True)).lower() == 'true' # global pin_memory for dataloaders # Get orientation exif tag for orientation in ExifTags.TAGS.keys(): @@ -83,7 +84,7 @@ def exif_transpose(image): 5: Image.TRANSPOSE, 6: Image.ROTATE_270, 7: Image.TRANSVERSE, - 8: Image.ROTATE_90,}.get(orientation) + 8: Image.ROTATE_90}.get(orientation) if method is not None: image = image.transpose(method) del exif[0x0112] @@ -144,7 +145,7 @@ def create_dataloader(path, shuffle=shuffle and sampler is None, num_workers=nw, sampler=sampler, - pin_memory=True, + pin_memory=PIN_MEMORY, collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn, worker_init_fn=seed_worker, generator=generator), dataset @@ -1152,6 +1153,6 @@ def create_classification_dataloader(path, shuffle=shuffle and sampler is None, num_workers=nw, sampler=sampler, - pin_memory=True, + pin_memory=PIN_MEMORY, worker_init_fn=seed_worker, generator=generator) # or DataLoader(persistent_workers=True)