pytorch如何使用cpu和gpu

第一、设置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)