[linux] pytorch 安装 torch torchvision gpu 配置

Start Locally | PyTorch （官网）

pytorch 安装

1. 查看cuda版本： 

cat /usr/local/cuda/version.txt

10.2.89

2. 官网查看下载命令

如果是镜像源下载，去掉-c pytorch。

Start Locally | PyTorch

conda install pytorch torchvision torchaudio cudatoolkit=10.2
#conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch

#pip install --upgrade torchvision

 版本： pytorch 1.7.1 torchvision 0.8.2  要求python>=3.6

 3. 测试

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

torch.cuda.is_available()
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')