torch-loader(example): use prefetch and try to run example in linux

iterative · Dec 11, 2024 · 59c7aa6 · 59c7aa6
1 parent 6ca3c98
commit 59c7aa6
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/examples/get_started/torch-loader.py b/examples/get_started/torch-loader.py
@@ -5,6 +5,7 @@
 
 """
 
+import multiprocessing
 import os
 from posixpath import basename
 
@@ -54,6 +55,7 @@ def forward(self, x):
 if __name__ == "__main__":
     ds = (
         DataChain.from_storage(STORAGE, type="image")
+        .settings(cache=True, prefetch=25)
         .filter(C("file.path").glob("*.jpg"))
         .map(
             label=lambda path: label_to_int(basename(path)[:3], CLASSES),
@@ -64,8 +66,9 @@ def forward(self, x):
 
     train_loader = DataLoader(
         ds.to_pytorch(transform=transform),
-        batch_size=16,
-        num_workers=2,
+        batch_size=25,
+        num_workers=4,
+        multiprocessing_context=multiprocessing.get_context("spawn"),
     )
 
     model = CNN()

diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py
@@ -6,14 +6,7 @@
 
 import pytest
 
-get_started_examples = sorted(
-    [
-        filename
-        for filename in glob.glob("examples/get_started/**/*.py", recursive=True)
-        # torch-loader will not finish within an hour on Linux runner
-        if "torch" not in filename or os.environ.get("RUNNER_OS") != "Linux"
-    ]
-)
+get_started_examples = sorted(glob.glob("examples/get_started/**/*.py", recursive=True))
 
 llm_and_nlp_examples = sorted(glob.glob("examples/llm_and_nlp/**/*.py", recursive=True))