记录使用GPU

CUDA with PyTorch

我们服务器显卡配置如下

CUDA版本是11.2，但是查看torch官网后发现没有cuda11.2版本对应的torch下载。
考虑到版本向下兼容，可以选择下载cuda11.1的版本。安装命令如下

# CUDA 11.1
pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu111/torch_stable.html

详情链接：https://pytorch.org/get-started/previous-versions/
注意：为了与CUDA11.1兼容，我们创建环境预设python版本是3.8.

另一个服务器的CUDA版本是11.4，可以下载cuda11.3的版本

# CUDA 11.3
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113

更新

由于我们的服务器驱动更新
现在可以使用最新版本的torch和pytorch_lightning命令如下

# CUDA 12.1
pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu121

解决122服务器的小bug

运行含有source命令时提示

source: not found

解决办法：运行如下指令即可

exec bash

参考链接：解决source命令bug

如果还有pytorch_lightning 需要安装

需要考虑pytorch_lightning 与python已经pytorch之间的版本兼容问题，为此需要先安装适合版本的pytorch-Lightning包，然后再按照上述方法安装torch。(注：pytorch-lightning和torch的安装都选pip，否则pytorch会将torch版本更新至最新版本)

pip install pytorch_lightning==1.9

pytorch-Lightning与torch之间的版本对应关系如下

详情可见网址: torch与pytorch_lightning版本依赖关系
驱动更新后，可以使用pytorch-lightning 2.2版本，安装命令如下

pip install pytorch_lightning==2.2

安装好后测试

import torch
 
print(torch.cuda.is_available()) # cuda是否可用
# true 
print(torch.cuda.current_device()) # 返回当前设备索引
# 0
print(torch.cuda.device_count())  # 返回GPU的数量
# 6
print(torch.cuda.get_device_name(0))