Smaller Python ML containers

What we saw

Python ML containers can become huge without anyone making an explicit decision. The common PyTorch install pulls CUDA libraries even for CPU-only workloads, which slows installs, bloats images, and can make serverless or autoscaled deploys slower.

Why the images were so big

The default package path optimizes for broad hardware support. If your workload does not need a GPU, the CUDA baggage is pure deployment weight.

# Before
pip install torch torchvision
# Can pull roughly 2.7 GB of CUDA-related packages

# After
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
# CPU-only path, roughly 200 MB in the source note

What we changed

Install CPU-only PyTorch before the rest of your requirements, and pin versions so a future upstream change does not alter your image unexpectedly.

RUN pip install --no-cache-dir torch==2.6.0+cpu \
    --extra-index-url https://download.pytorch.org/whl/cpu && \
    pip install --no-cache-dir -r requirements.txt

Combine that with normal Docker hygiene: slim base images, fewer layers, same-layer cleanup for apt caches, --no-install-recommends, and pip --no-cache-dir.

What it did to the numbers

What we would tell another team

Container size is often dependency selection disguised as infrastructure. If a workload is CPU-only, make that choice explicit in the Dockerfile.

@article{gupta2025smaller,
  author = {Lalit Gupta},
  title = {Smaller Python ML containers},
  journal = {VideoDB Labs},
  year = {2025},
  month = {aug},
  note = {https://videodb.io/blog/python-ml-container-size},
}