While trying to reduce the size of a Docker image, I noticed pip install torch adds a few GB. A big chunk of this comes from [...]/site-packages/nvidia. Since I'm not using a GPU, I'd like to not install the nvidia things.
Here is a minimal example:
FROM python:3.12.5
RUN pip install torch
(Ignoring -slim base images, since this is not the point here.)
Resulting size:
FROM python:3.12.5 -> 1.02GBRUN pip install torch -> 8.98GBRUN pip install torch && pip freeze | grep nvidia | xargs pip uninstall -y instead -> 6.19GB.While the last point reduces the final size, all the nvidia stuff is still downloaded and installed, which costs time and bandwidth.
So, how can I install torch without nvidia directly?
Using --no-deps is not a convenient solution, because of the other transitive dependencies, that I would like to install.
Of course, I could explicitly list every single one, but looking at this list of packages installed with torch
mpmath
typing-extensions
sympy
nvidia-nvtx-cu12
nvidia-nvjitlink-cu12
nvidia-nccl-cu12
nvidia-curand-cu12
nvidia-cufft-cu12
nvidia-cuda-runtime-cu12
nvidia-cuda-nvrtc-cu12
nvidia-cuda-cupti-cu12
nvidia-cublas-cu12
networkx
MarkupSafe
fsspec
filelock
triton
nvidia-cusparse-cu12
nvidia-cudnn-cu12
jinja2
nvidia-cusolver-cu12
torch
I'd like to avoid manually maintaining this list since it would change with future versions of torch.
As (roundaboutly) documented on pytorch.org's getting started page, Torch on PyPI is Nvidia enabled; use the download.pytorch.org index for CPU-only wheels:
RUN pip install torch --index-url https://download.pytorch.org/whl/cpu
Also please remember to specify a somewhat locked version of Torch, e.g.
RUN pip install torch~=2.4.0 --index-url https://download.pytorch.org/whl/cpu