requirements.txt for pytorch for both CPU and GPU platforms

user_2713740 0
Viewed 341

I am trying to create a requirements.txt to use pytorch but would like it to work on both GPU and non-GPU platforms.

I do something like on my Linux GPU system:

--find-links https://download.pytorch.org/whl/cu113/torch_stable.html

torch==1.10.2+cu113
torchvision==0.11.3+cu113
pytorch-lightning==1.5.10

This works fine and the packages are installed and I can use the GPU-enabled pytorch.

I wonder how I can modify this for the mac and non GPU users to install the non cuda package for torch and torchvision? Do I need to maintain separate requirements.txt files?

requirements.txt pip python
1 Answers

Check https://pytorch.org/. You will see that "MacOS Binaries dont support CUDA, install from source if CUDA is needed". However, you can still get performance boosts (this will depend on your hardware) by installing the MPS accelerated version of pytorch by:

# MPS acceleration is available on MacOS 12.3+
pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

This command can be generated here: https://pytorch.org/

enter image description here

And in order to install different Torch versions on different platorms you can use conditionals in your requirements.txt like this

# for CUDA 11.3 torch on Linux
--find-links https://download.pytorch.org/whl/cu113/torch_stable.html; sys_platform == "linux"
torch==1.10.2; sys_platform == "linux"
torchvision==0.11.3; sys_platform == "linux"
pytorch-lightning==1.5.10; sys_platform == "linux"

# for MPS accelerated torch on Mac
--pre --extra-index-url https://download.pytorch.org/whl/nightly/cpu; sys_platform == "darwin"
torch==1.10.2; sys_platform == "darwin"
torchvision==0.11.3; sys_platform == "darwin"
pytorch-lightning==1.5.10; sys_platform == "darwin"

# for CPU torch on Mac
# torch==1.10.2; sys_platform == "darwin"
# torchvision==0.11.3; sys_platform == "darwin"
# pytorch-lightning==1.5.10; sys_platform == "darwin"

This will install CUDA enabled torch and torchvision on Linux but the MPS accelerated version of them on MacOS

user_95067730
Related