Finding the Bottleneck

Many people pay for a powerful G P U, and then they let it sit idle, because it's waiting for data. A simple check can double your training speed for free. And you won't need to change your model at all.

• Expensive GPU, sitting idle
• Waiting for data
• A free speed-up

Timing the two halves

Training has two main parts. Getting data, and computing. We time each one separately. Run it. Now you can see which half is slowing you down. If loading takes longer, your G P U is waiting. The fix is more worker processes in the data loader, faster storage, or caching preprocessed data. Tools such as N V I D I A S M I show you G P U use. If it's low, then your data is the problem.

import time
import torch

x = torch.rand(256, 3, 64, 64)
conv = torch.nn.Conv2d(3, 64, 3, padding=1)

start = time.time()
for _ in range(10):
    loaded = x.clone()
load = time.time() - start

start = time.time()
for _ in range(10):
    y = conv(x)
compute = time.time() - start

print(f"Loading: {load:.3f}s   Computing: {compute:.3f}s")
print("If loading dominates, your GPU is sitting idle")

Loading: 0.003s   Computing: 0.687s
If loading dominates, your GPU is sitting idle

Quick wins

Here are the quick wins. Set num workers in your data loader, to use more processes. Use pin memory for faster transfers. Keep images on a fast disk. Use a larger batch, if your memory allows it. And use mixed precision. These simple settings can make training several times quicker, with no change to the model.

Summary