From d6c7900d06c3388bf814cecbe90f91a9afecbefb Mon Sep 17 00:00:00 2001 From: David Wade Date: Fri, 27 Jul 2018 17:34:54 +0200 Subject: [PATCH] Use scandir in _find_classes (#557) (#559) * Use scandir in _find_classes (#557) * fix lint --- 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..3bd4c485b65 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