Repository: lightly-ai/lightly

Issue:
Lightly Dataset is slow at indexing larger datasets
`LightlyDataset` is significantly slower than a simple `os.listdir` (500ms vs 50ms with a folder of 100k images). Let's figure out whether we can improve.

**Tasks:**

- [x] Figure out where we spend most of the time
- [x] Figure out whether we can improve


Discussion:
I looked briefly into this. There are a few points where we could get 20-30% more speed. But they would require us to use a custom dataset instead of relying on the torchvision `ImageFolder`. 

I made an issue here: https://github.com/pytorch/vision/issues/3902
This is related to [617](https://github.com/lightly-ai/lightly-core/issues/617).

Some of the reasons:
- `LightlyDataset` makes several calls to `os.listdir` (total of 3). This could be done more efficiently.
- `LightlyDataset` uses `os.path.join` which can be slow if done repeatedly but there is no way around this.

Proposal: Let's be more resourceful with the calls to `os.listdir` in [this function](https://github.com/lightly-ai/lightly/blob/ffabd63bd340c1cc69def46b3c772dcfc783efa4/lightly/data/_helpers.py#L71). This should carve off an easy 100ms on the dataset you described above.

Bonus: If there's a quick way to figure out whether a directory contains subdirectories which doesn't require `os.path.join` (as opposed to [our function](https://github.com/lightly-ai/lightly/blob/ffabd63bd340c1cc69def46b3c772dcfc783efa4/lightly/data/_helpers.py#L54)), we could further improve the performance.
Apparently, `os.scandir` can also lead to signficant speed-ups. See:
- https://stackoverflow.com/questions/800197/how-to-get-all-of-the-immediate-subdirectories-in-python
- https://docs.python.org/3/library/os.html#os.scandir

Resolving patch:
diff --git a/lightly/data/_helpers.py b/lightly/data/_helpers.py
--- a/lightly/data/_helpers.py
+++ b/lightly/data/_helpers.py
@@ -32,10 +32,8 @@ def _contains_videos(root: str, extensions: tuple):
     Returns:
         True if root contains subdirectories else false.
     """
-    list_dir = os.listdir(root)
-    is_video = \
-        [f.lower().endswith(extensions) for f in list_dir]
-    return any(is_video)
+    with os.scandir(root) as scan_dir:
+        return any(f.name.lower().endswith(extensions) for f in scan_dir)
 
 
 def _is_lightly_output_dir(dirname: str):
@@ -61,11 +59,9 @@ def _contains_subdirs(root: str):
         True if root contains subdirectories else false.
 
     """
-    list_dir = os.listdir(root)
-    list_dir = list(filter(lambda x: not _is_lightly_output_dir(x), list_dir))
-    is_dir = \
-        [os.path.isdir(os.path.join(root, f)) for f in list_dir]
-    return any(is_dir)
+    with os.scandir(root) as scan_dir:
+        return any(not _is_lightly_output_dir(f.name) for f in scan_dir \
+            if f.is_dir())
 
 
 def _load_dataset_from_folder(root: str, transform):
diff --git a/lightly/data/_image.py b/lightly/data/_image.py
--- a/lightly/data/_image.py
+++ b/lightly/data/_image.py
@@ -3,13 +3,15 @@
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
 
+from typing import List, Tuple
+
 import os
 import torchvision.datasets as datasets
 
 from lightly.data._image_loaders import default_loader
 
 
-def _make_dataset(directory, extensions=None, is_valid_file=None):
+def _make_dataset(directory, extensions=None, is_valid_file=None) -> List[Tuple[str, int]]:
     """Returns a list of all image files with targets in the directory.
 
     Args:
@@ -41,16 +43,17 @@ def _is_valid_file(filename):
         _is_valid_file = is_valid_file
 
     instances = []
-    for fname in os.listdir(directory):
+    for f in os.scandir(directory):
 
-        if not _is_valid_file(fname):
+        if not _is_valid_file(f.name):
             continue
 
-        path = os.path.join(directory, fname)
-        item = (path, 0)
+        # convention: the label of all images is 0, based on the fact that
+        # they are all in the same directory
+        item = (f.path, 0)
         instances.append(item)
 
-    return instances
+    return sorted(instances, key=lambda x: x[0]) # sort by path
 
 
 class DatasetFolder(datasets.VisionDataset):


Test patch: