今天在测试stablediffusion的时候,使用GPU进行加速,需要安装CUDA和Pytorch,期间遇到一些坑记录下来,期望同学们不要和我一样再一次踩坑。

一、安装配套的CUDA驱动和pytorch

  • 在cmd窗口执行查询命令nvidia-smi,找到自己显卡的cuda版本
    在这里插入图片描述

  • 到nvidia官网下载cuda安装包,版本向下兼容,可以下载比显卡cuda版本低的驱动 https://developer.nvidia.com/cuda-toolkit-archive 下载前先看pytorch可以适配使用的版本
    在这里插入图片描述

  • 在这里先不要直接下载,先到pytorch官网上找到对应的cuda驱动的版本,pytorch官网地址:https://pytorch.org/get-started/previous-versions/
    比如这里适配最新的cuda版本是12.1,那我们就下载cuda 12.1版本的驱动
    在这里插入图片描述

  • 驱动下载完成后直接按照提示精简安装就可以

  • 在python环境中使用pip安装pytorch,这里一定要选择带cuda注释的,下错了就用的不是GPU了,也要对应cuda驱动的版本
    在这里插入图片描述

CUDA 11.8
pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu118
CUDA 12.1
pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121
  • 如果是已经安装了cuda和pytorch,cuda驱动可以覆盖安装,pytorch需要先使用pip先卸载,再进行安装

二、安装完成进行测试

  • 在python控制台中输入如下代码进行测试
import torch
# 查看pytorch版本号,cu结尾表示GPU使用的,cpu结尾表示cpu使用的
print(torch.__version__)

# 查看cuda驱动版本好
print(torch.version.cuda)

# 查看cuda驱动版本和pytorch版本是否匹配
print(torch.cuda.is_available())

在这里插入图片描述

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐