From 0613005ae3a730f467592d41756c386ce997eb67 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Wed, 9 Jul 2025 21:50:20 +0200 Subject: [PATCH] Update main.py tutorial, set persistent_workers, make GPU execution much faster Since we are setting num_workers to 1, we should either make this worker persistent or not set the num_workers. On my machine, running this script as is takes 2.21 mins, setting the workers to persistent or removing num_workers (setting to its default of 0) makes it a good deal faster at 1.44 and 1.66 mins respectively. I would opt for adding the persistent_workers setting as it is a little faster and makes it trivial for even a beginner to speed it up further by increasing the num_workers (setting to 10 reduced my run time to 0.53 mins --- mnist/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mnist/main.py b/mnist/main.py index 0671b8e336..dee5a384cb 100644 --- a/mnist/main.py +++ b/mnist/main.py @@ -107,6 +107,7 @@ def main(): test_kwargs = {'batch_size': args.test_batch_size} if use_accel: accel_kwargs = {'num_workers': 1, + 'persistent_workers': True, 'pin_memory': True, 'shuffle': True} train_kwargs.update(accel_kwargs)