Search code examples
pythonpiprequirements.txt

requirements.txt for pytorch for both CPU and GPU platforms

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?

Solution

  • February 2024 update

    Check https://pytorch.org/. You will see that "CUDA is not available on MacOS, please use default package". 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 torch torchvision torchaudio

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

    enter image description here

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

    # for CUDA 11.8 torch on Linux
--index-url https://download.pytorch.org/whl/cu118; sys_platform == "linux"
torch; sys_platform == "linux"
torchvision; sys_platform == "linux"
pytorch-lightning; sys_platform == "linux"

# for MPS accelerated torch on Mac
torch; sys_platform == "darwin"
torchvision; sys_platform == "darwin"
pytorch-lightning; sys_platform == "darwin"

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