第一、设置cpu和gpu
-
判定是否存储cuda:use_cuda = torch.cuda.is_available() and not args.no_cuda
-
选择设备参数cuda为gpu,cpu为采用cpu运行 device = torch.device('cuda' if use_cuda else 'cpu')
第二、cudnn的选择方式
由于cpu或者低版本的gpu对于pytorch会保存,错误一般为:CUDNN_STATUS_ARCH_MISMATCH
设置不使用cudnn库:
- torch.backends.cudnn.enabled = False
第三、gpu的设备选择
- 利用环境变量设置
import os
os.environ["CUDA_VISIBLE_DEVICES"] = 0
- 命令行设置
CUDA_VISIBLE_DEVICES = 0 python train.py
- 代码行设置
torch.cuda.set_device(6)