Ubuntu22.04下使用Conda安装PyTorch GPU版本

环境

首先，你需要有 GPU 支持。
Ubuntu 22.04
显卡

$ nvidia-smi
Mon May 22 11:15:33 2023
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 530.30.02              Driver Version: 530.30.02    CUDA Version: 12.1     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                  Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf            Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA A100 80GB PCIe           On | 00000000:21:00.0 Off |                    0 |
| N/A   34C    P0               45W / 300W|      4MiB / 81920MiB |      0%      Default |
|                                         |                      |             Disabled |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      3309      G   /usr/lib/xorg/Xorg                            4MiB |
+---------------------------------------------------------------------------------------+

安装

Conda 安装

使用 Conda 的最大好处就是包环境管理。最大坏处就是屎一样大。

下载

目前 Conda 最新版本是 2023.03-1，参考 https://repo.anaconda.com/archive/。
使用命令行

wget https://repo.anaconda.com/archive/Anaconda3-2023.03-1-Linux-x86_64.sh

安装

使用命令行

chmod +x Anaconda3-2023.03-1-Linux-x86_64.sh
./Anaconda3-2023.03-1-Linux-x86_64.sh

即可完成安装。

Conda 常用命令

查看虚拟环境

使用命令行

conda env list

创建虚拟环境

使用命令行

conda create --name 虚拟环境名字 python=3.8

python=3.8 是指定 python 软件包版本。

激活虚拟环境

使用命令行

conda activate 虚拟环境名字

退出虚拟环境

使用命令行

conda deactivate

删除虚拟环境

使用命令行

conda remove -n 虚拟环境名字 --all

PyTorch 官网支持

请到官网 https://pytorch.org/get-started/locally/，Torch 提供了对应的选择。选择如下

创建 torch 虚拟环境

conda create --name torch python=3.8

激活 torch 虚拟环境

conda activate torch

激活成功后，将出现以下的状态

(torch) 用户名@服务器名:~

安装 torch

参考官网给出的命令行。

conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

我前面几次安装 torch 都是 cpu 版本。惭愧参看了好多资料才正确安装 GPU。否则浪费了 A100 最新版本。

验证安装

使用 python 进入命令行状态

Python 3.8.16 (default, Mar  2 2023, 03:21:46)
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.__version__) #打印torch版本
2.0.1
>>> print(torch.backends.cudnn.version()) #打印cudnn版本
8700
>>> print(torch.version.cuda) #打印cuda版本
11.8
>>> torch.cuda.is_available() #cuda是否使用
True
>>> print(torch.cuda.get_device_name(0)) #打印GPU
NVIDIA A100 80GB PCIe

这样，我们就打完收工。