使用Poetry安装PyTorch

一、环境说明

python：3.10
操作系统：windows 11
pytorch：支持gpu的版本

二、确认操作系统支持的cuda版本

使用win+R打开命令提示行，并输入nvidia-smi

可以看到右上角有个CUDA Version: 12.8，注意，这里不是你现在电脑的CUDA版本，而是你现在nvidia驱动所支持的最高CUDA版本。

三、选择PyTorch版本

既然要用，那就直接用gpu版本的，而不是cpu版本的

3.1. 一份可用的PyTorch依赖

我知道有些同学并不想要解题过程，而是亟需一份立刻能用的。那么，这里分享我现在用的一份依赖清单，pyproject.toml文件如下，复制这里的内容到本地，然后poetry install就可以安装依赖了。

[tool.poetry]
name = "elf-pytorch"
version = "0.1.0"
description = ""
authors = ["TreeOfWorld"]
readme = "README.md"

[tool.poetry.dependencies]
python = ">=3.10,<3.12.0"
torch = { version = "2.6.0", extras = ["cu124"], source = "torch" }
torchvision = { version = "0.21.0", extras = ["cu124"], source = "torch" }
torchaudio = { version = "2.6.0", extras = ["cu124"], source = "torch" }
numpy = "^2.2.4"

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cu124"
priority = "explicit"

[[tool.poetry.source]]
name = "aliyun"
url = "http://mirrors.aliyun.com/pypi/simple/"
priority = "primary"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"

3.2. 验证环境安装成功

运行以下代码来验证环境安装是否成功

import torch


def test_gpu():
    # 测试cuda是否可用
    print(f"Support CUDA? :{torch.cuda.is_available()}")
    x = torch.tensor([10.0])
    x = x.cuda()

    y = torch.randn(2, 3)
    y = y.cuda()

    z = x + y
    print(z)

    # 测试cudnn是否可用
    from torch.backends import cudnn
    print(f"Support cudnn? :{cudnn.is_available()}")


print(torch.__version__)

if __name__ == '__main__':
    test_gpu()

若运行结果如下，则说明pytorch安装成功，且cuda和cudnn是可用的

3.3. 选择pytorch版本的方法

这一段是记录我选择上述pytorch版本的过程

3.3.1. 确认操作系统支持的cuda版本

在步骤二中知道，我当前操作系统支持的cuda版本最高为12.8，

3.3.2. 在pytorch官方源找一套可以用的版本

pytorch官方源地址：https://download.pytorch.org/whl
我们需要依次确认torch、torchvision、torchaudio的版本

确认torch版本

首先在官方源地址搜索torch

点击链接进入torch的版本列表，查看可以下载的所有版本，大概可以看到当前torch最新的版本为2.6.0，因为要安装gpu版本的，所以要在页面搜索：2.6.0+cu

可以看到，2.6.0版本只支持cu118、cu124、cu126三个版本，因为我们的操作系统可以安装cuda12.8及其以下的版本，所以以上的三个版本都是可选的。

确认torchvision版本

同样的方法，我们在torchvision所在的页面，查看有哪些版本可选，略微浏览就可以看出，最新的版本为0.21.0，同样因为要安装gpu版本的，所以我们搜索0.21.0+cu，可以看到torchvision也是支持cu118、cu124、cu126三个版本

确认torchaudio版本

同样的方法，看到torchaudio的版本最新为2.6.0，搜索2.6.0+cu后发现，可以看到torchaudio也是支持cu118、cu124、cu126三个版本

3.3.3. 编写pyproject.toml文件

以下只列出pytorch相关的内容，完整文件的格式可参考3.1中所给出的，以下列出三个版本所使用的配置方式，主要差异在tool.poetry.dependencies.extras和tool.poetry.source.url
如果选择使用cuda12.4，那么配置如下：

......
[tool.poetry.dependencies]
torch = { version = "2.6.0", extras = ["cu124"], source = "torch" }
torchvision = { version = "0.21.0", extras = ["cu124"], source = "torch" }
torchaudio = { version = "2.6.0", extras = ["cu124"], source = "torch" }

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cu124"
priority = "explicit"
......

如果选择使用cuda12.6，那么配置如下：

......
[tool.poetry.dependencies]
torch = { version = "2.6.0", extras = ["cu126"], source = "torch" }
torchvision = { version = "0.21.0", extras = ["cu126"], source = "torch" }
torchaudio = { version = "2.6.0", extras = ["cu126"], source = "torch" }

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cu126"
priority = "explicit"
......

如果选择使用cuda11.8，那么配置如下：

......
[tool.poetry.dependencies]
torch = { version = "2.6.0", extras = ["cu118"], source = "torch" }
torchvision = { version = "0.21.0", extras = ["cu118"], source = "torch" }
torchaudio = { version = "2.6.0", extras = ["cu118"], source = "torch" }

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cu118"
priority = "explicit"
......

需要注意的是，因为pytorch依赖于numpy，如果选择cuda11.8的话，numpy需要使用1.x.x+的版本
使用cuda12.4或者cuda12.6的话，numpy可以使用2.x.x+的版本

原因：随着大版本从1升到2，numpy的api发生了变化