python - CPU和GPU平台的pytorch的requirements.txt

Question

我正在尝试创建一个 requirements.txt 来使用 pytorch，但希望它能够在 GPU 和非 GPU 平台上工作。

我在我的 Linux GPU 系统上做了类似的事情:

--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

这工作正常，软件包已安装，我可以使用支持 GPU 的 pytorch。

我想知道如何修改此设置，以便 Mac 和非 GPU 用户安装 torch 和 torchvision 的非 cuda 软件包？我需要维护单独的 requirements.txt 文件吗？

Answer 1

检查https://pytorch.org/ 。您将看到“MacOS 二进制文件不支持 CUDA，如果需要 CUDA，请从源代码安装”。但是，您仍然可以通过安装 MPS 加速版本的 pytorch 来获得性能提升(这取决于您的硬件):

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

可以在此处生成此命令:https://pytorch.org/

为了在不同的平台上安装不同的 Torch 版本，您可以在 requirements.txt 中使用条件，如下所示

# 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"

这将在 Linux 上安装支持 CUDA 的 torch 和 torchvision，但在 MacOS 上安装它们的 MPS 加速版本