Pytorch-lightning CPU-only installation

[sivakhno]

Hello all - just wanted to discuss a use-case with CPU vs GPU PL install. We do normal training on GPUs, but when deploying for prediction we use CPUs and would like to keep the Docker container size as small as possible. It't not clear if it's possible to install pytorch-lightning with CPU-only torch distribution, which is much smaller. Is there any possible equivalent to pip install pytorch-lightning[cpu] . Thanks for suggestions!

[awaelchli]

Hey, if you install pytorch first (cpu only) and then Lightning it will just use that version.

[cverrier]

Installing the CPU-only version of PyTorch first is a good solution in the case one does not rely on fully-locked requirements files, such as those generated by pip-compile.

In the case one wants to use these fully-locked requirements files, installing the CPU-only version of PyTorch first does not work anymore.

One workaround could be to manually edit the fully-locked requirements files (generated with e.g. pip-compile), but this could be cumbersome because of eventual transitive dependencies (e.g. PyTorch depends on several NVIDIA libraries).

It would be nice to be able to download a CPU-only version of Lightning directly.

[mbrzus]

Hi @cverrier, were you able to find the solution to that?

[cverrier]

Hi @cverrier, were you able to find the solution to that?

I could only think about a workaround that uses pip install's --index-url and --extra-index-url flags. You can create a requirements.in file, and write the following inside it:

1. Change the default index URL and install PyTorch's CPU-only version:

   --index-url https://download.pytorch.org/whl/cpu
   torch
   torchvision
   torchaudio
   
2. Specify an extra index URL that points to the default one (i.e. PyPI) and install Lightning:

   --extra-index-url https://pypi.org/simple
   lightning
   

Then generating the fully-locked requirements should be straightforward, by simply running

pip-compile requirements.in

In terms of dependency injection attacks, I think one should be safe because the default index URL now points to one owned by PyTorch, which is unlikely to contain suspicious Python dependencies.