Skip to content

Commit

Permalink
Use scandir in _find_classes (#557) (#559)
Browse files Browse the repository at this point in the history
* Use scandir in _find_classes (#557)

* fix lint
  • Loading branch information
David Wade authored and soumith committed Jul 27, 2018
1 parent f1b5907 commit d6c7900
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion torchvision/datasets/folder.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import os
import os.path
import sys


def has_file_allowed_extension(filename, extensions):
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit d6c7900

Please sign in to comment.