From 3146b18f9196b62ddcb42908d1bc235b024c4124 Mon Sep 17 00:00:00 2001 From: David Wade Date: Fri, 27 Jul 2018 11:43:37 +0000 Subject: [PATCH 1/2] Use scandir in _find_classes (#557) --- torchvision/datasets/folder.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py index 52362ca198f..23f0ef7d482 100644 --- a/torchvision/datasets/folder.py +++ b/torchvision/datasets/folder.py @@ -4,6 +4,7 @@ import os import os.path +import sys def has_file_allowed_extension(filename, extensions): @@ -110,7 +111,11 @@ def _find_classes(self, dir): Ensures: No class is a subdirectory of another. """ - classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))] + if sys.version_info >= (3,5): + # Faster and available in Python 3.5 and above + classes = [d.name for d in os.scandir(dir) if d.is_dir()] + else: + classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))] classes.sort() class_to_idx = {classes[i]: i for i in range(len(classes))} return classes, class_to_idx From 1192ccdc804661eb0dc8855706de56b56d078777 Mon Sep 17 00:00:00 2001 From: Soumith Chintala Date: Fri, 27 Jul 2018 10:40:58 -0400 Subject: [PATCH 2/2] fix lint --- torchvision/datasets/folder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py index 23f0ef7d482..3bd4c485b65 100644 --- a/torchvision/datasets/folder.py +++ b/torchvision/datasets/folder.py @@ -111,7 +111,7 @@ def _find_classes(self, dir): Ensures: No class is a subdirectory of another. """ - if sys.version_info >= (3,5): + if sys.version_info >= (3, 5): # Faster and available in Python 3.5 and above classes = [d.name for d in os.scandir(dir) if d.is_dir()] else: