快速构建基于GPU-CUDA的PyTorch

我目前使用的是一个基于RTX 3090的深度学习系统，3090显卡的CUDA版本为11.7，如果要使用GPU训练模型就只能自行构建PyTorch。

背景

显卡 RTX 3090 24G

CUDA版本 11.7

Cudnn版本 v8

OS Ubuntu 20.04

Python 3.8

构建流程

安装依赖

# 通用依赖
conda install astunparse numpy ninja pyyaml setuptools cmake cffi typing_extensions future six requests dataclasses

# Linux依赖
conda install mkl mkl-include

# CUDA依赖
conda install -c pytorch magma-cuda117 # 或者 magma-cuda* *号与 https://anaconda.org/pytorch/repo中你的CUDA版本匹配的数字

获取PyTorch源码

git clone -b v1.12.1 --recursive https://github.com/pytorch/pytorch
cd pytorch
# 如果更新一个已存在的检出：
git submodule sync
git submodule update --init --recursive --jobs 0

安装PyTorch

export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
python setup.py install

# 建议使用Python版本3.7.6+ 或者3.8.1+

等待构建完成，即可使用，构建根据机器配置不同用时会有差异，不过我用了挺久构建的。

检验PyTorch以及CUDA支持

可以使用PyCharm的Python Console或者在终端中输入 python 进入控制台

>>> import torch

>>> torch.cuda.is_available()
True

>>> torch.cuda.device_count()
4

>>> torch.cuda.current_device()
0

>>> torch.cuda.device(0)
<torch.cuda.device at 0x7efce0b03be0>

>>> torch.version.cuda

参考文献

GitHub：pytorch / pytorch