python 安装 cuda 版 pytorch

菜猿
2,653 阅读4分钟

注意：若之前已经使用 pip install pytorch 等类似命令直接安装，需要先卸载后再安装

一. 查看设备中安装的CUDA版本

执行命令 nvidia-smi 在输出结果中 CUDA Version: 位置可找到CUDA版本，如下图，版本为 12.6

image.png

二. 查找版本并安装

（1）查询对应版本

打开PyTorch 官网，下划找到 Install PyTorch 部分，根据CUDA Version、操作系统、包类型等条件选择

注意： 选择的 CUDA Version 只能低于设备中安装的实际版本

  • 例如在第一步中获取到的CUDA版本为 12.6，此时选择的版本必须小于等于12.6，即 12.4

image.png

最终 Run this Command: 部分即为安装pytorch的命令

历史版本

以上位置仅能查询到最新版本，若想安装其他版本，可进入Previous PyTorch Versions | PyTorch页面，或点击官网中以下按钮

image.png

在新页面中即可找到使用 conda 命令安装历史版本的命令

image.png

（2）安装

方式一、直接运行命令安装

复制上一步获得的命令在Terminal中执行即可.

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

方式二、手动下载wheel包

1. 明确包版本号含义

  • CUDA 版本

    • 根据实际的CUDA版本号得到wheel包中的版本号非常简单，把句点去掉再添加前缀即可
    • 例如设备中CUDA版本 12.6 对应 126，再添加说代表CUDA的前缀cu，得到 cu126

    再次说明，对于CUDA版本，只能低于设备中安装的实际版本
    在实际寻找版本时需要寻找小于等于126的版本
    如 cu124 cu121 等

  • pytorch 相关包版本

    • 版本号分为两部分，例如 2.5.1+cu124
      • 第一部分为包本身的版本，即 2.5.1 部分
      • 第二部分为CUDA版本，即 cu124 部分
      • 中间以加号 + 连接
      • 最终 2.5.1+cu124 表示torch是运行在CUDA12.4基础上的、本身版本为2.5.1的版本
2. 查询不同包之间版本对应关系

最简单的方式是在 Previous PyTorch Versions | PyTorch 页面给出的安装命令中获取不同包版本的对应关系，若所需包不存在可使用以下方式

一般可直接在不同包的GitHub项目页找到对应关系，可先进入 PyTorch组织主页，如下图

image.png

找到之后点击进入项目主页，找到安装说明，一般是 Installation 部分，会有两种情况

  • 以 torchvision 为例，可直接在项目主页找到版本对应关系
    实际版本号为三段式，程序员应该很熟悉，即 大版本.小版本.补丁版本，此处只需要关注前两段即 大版本.小版本 部分
    如图中torch中 2.5 版本可以代表 2.5.1 2.5.2 2.5.3 ... 2.5.111 等所有版本
    其他包含义一致 image.png

  • 以 torchaudio 为例，项目主页中没有直接提供版本对应关系，而是给出了官方文档地址 image.png
    点击进入，如下图，滑动页面找到版本对应关系部分，版本号含义同上 image.png

真实的版本可直接在GitHub release 中找到，不做赘述

2. 根据对应关系下载

打开PyTorch源站（其实第一步中获取到安装命令也是在这个地址中下载的），如下图

image.png

可在其中找到 torch torchaudio torchvision 等包，点击跳转进入

以 torch 包为例，跳转进入以下页面

假如设备中实际的CUDA版本号为 12.6，想要安装2.5.1版本的torch
由于torch版本是可以直接在GitHub release 中获取到的，所以torch版本是精确的
那么可以搜索 -2.5.1+cu12，如下图

image.png

可发现找不到 cu126 的cuda版本，因此可在126基础上向下找，找到了 cu124 版本

  • python版本如cp310-cp310表示支持python3.10.x版本
  • 操作系统如 linux 表示支持Linux系统， win 表示支持Windows系统
    • x86_64 以及 amd64 均表示64位操作系统，以及x86架构（如Intel、AMD的CPU）
    • aarch64 表示64位操作系统，以及arm架构

找到后进行下载

3. 安装

使用 pip 安装上一步下载到的.whl 包文件即可，例如

pip install torch-2.5.1+cu124-cp310-cp310-win_amd64.whl

三. 验证

在 Python 命令行中执行一下python代码，若返回 True 则说明CUDA版python安装成功

import torch

torch.cuda.is_available()
avatar
文章
阅读
粉丝