Pytorch GPU 训练环境搭建

之前⼀直使⽤ Tensorflow 训练模型，第⼀次训练Pytorch模型的时候，发现速度很慢，仔细观察，发现GPU 内存占⽤为0，基本没有使⽤GPU。

AssertionError: CUDA unavailable, invalid device 0 requested
cuda不可⽤报错，现实没有有效的驱动可使⽤

测试cuda是否配置正确

import torch
print(torch.cuda.is_available())

重新安装cuda

检测本地GPU CUDA版本 nvidia-smi

正在上传…重新上传取消

1080×566 89.1 KB

pip3 install torch1.9.0+cu101 torchvision0.10.0+cu101 torchaudio=0.9.0 -f
https://download.pytorch.org/whl/torch_stable.html

当前 torchvision0.10.0+cu101 版本必须是匹配的。如果版本不匹配ÿ