Search code examples
pythonpytorchpipdockerfiletorch

How to install torch without nvidia?

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.02GB
  • After RUN pip install torch -> 8.98GB
  • With RUN 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.

Solution

  • 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