diff --git a/examples/get_started/torch-loader.py b/examples/get_started/torch-loader.py index 19ca484ca..a77e30257 100644 --- a/examples/get_started/torch-loader.py +++ b/examples/get_started/torch-loader.py @@ -19,7 +19,7 @@ from datachain.torch import label_to_int STORAGE = "gs://datachain-demo/dogs-and-cats/" -NUM_EPOCHS = os.getenv("NUM_EPOCHS", "3") +NUM_EPOCHS = int(os.getenv("NUM_EPOCHS", "3")) # Define transformation for data preprocessing transform = v2.Compose( @@ -68,7 +68,8 @@ def forward(self, x): train_loader = DataLoader( ds.to_pytorch(transform=transform), batch_size=25, - num_workers=4, + num_workers=max(4, os.cpu_count() or 2), + persistent_workers=True, multiprocessing_context=multiprocessing.get_context("spawn"), ) @@ -77,7 +78,7 @@ def forward(self, x): optimizer = optim.Adam(model.parameters(), lr=0.001) # Train the model - for epoch in range(int(NUM_EPOCHS)): + for epoch in range(NUM_EPOCHS): with tqdm( train_loader, desc=f"epoch {epoch + 1}/{NUM_EPOCHS}", unit="batch" ) as loader: diff --git a/noxfile.py b/noxfile.py index 72b8d9722..edb7c30f5 100644 --- a/noxfile.py +++ b/noxfile.py @@ -81,6 +81,8 @@ def examples(session: nox.Session) -> None: session.install(".[examples]") session.run( "pytest", + "--durations=0", + "tests/examples", "-m", "examples", *session.posargs, diff --git a/tests/examples/test_examples.py b/tests/examples/test_examples.py index d3735f364..6191b9593 100644 --- a/tests/examples/test_examples.py +++ b/tests/examples/test_examples.py @@ -52,7 +52,7 @@ def smoke_test(example: str, env: Optional[dict] = None): @pytest.mark.get_started @pytest.mark.parametrize("example", get_started_examples) def test_get_started_examples(example): - smoke_test(example, {"NUM_EPOCHS": "1"}) + smoke_test(example) @pytest.mark.examples