Pytorch + CUDA + GPU 的环境搭建

环境依赖一览

• Python 3.9.12
  3 版本基本均可
• NVDIA 517 显卡驱动
  选择自己合适的驱动，显卡 <–适配– 驱动 <–适配– CUDA，显卡版本可以在官网搜索查看，CUDA 和驱动的具体关系可以去官网文档 CUDA Documentation 查看
• CUDA 11.6
  选择自己适合的 CUDA 并下载安装。务必参考上句中的适配性，否则无法正常使用
• Anaconda 22.9.0
  用来便捷安装 pytorch-GPU 和 cudatoolkit 的，pip 由于其非python库的原因无法直接安装 toolkit（但实际上可以用官网所示代码用 --extra-index-url 来安装）
• Pytorch
  使用官网对应命令安装即可，比如 conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c nvidia -c conda-forge

安装细节

我们需要用镜像源加速安装，同时有些下载需要添加额外的源，这里我使用的是清华源

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/win-64

安装时注意使用管理员打开 cmd 窗口，不然可能导致最后报无法写入文件的错，如下所示

EnvironmentNotWritableError: The current user does not have write permissions to the target environment.
  environment location: D:\Anaconda3

安装包很大，要等待一会儿。

完毕后可以在 python 中检查是否可以使用 CUDA 与 GPU 运算

import torch

if __name__ == "__main__":
    print(torch.cuda.is_available())
    print(torch.cuda.device_count())
    print(torch.cuda.get_device_name(0))
    print(torch.cuda.current_device())

    module = torch.nn.Linear(3, 1)
    print(list(module.parameters())[0].device)
    module.cuda(0)
    print(list(module.parameters())[0].device)
    module.cpu()
    print(list(module.parameters())[0].device)

我的输出结果为

True
1
NVIDIA GeForce RTX 2060
0
cpu
cuda:0
cpu