PyTorch指定GPU的方法

1. 系统变量方式

（1）export CUDA_VISIBLE_DEVICES=0 #这里是要使用的GPU编号，正常的话是从0开始

（2）在程序开头设置os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'

（3）CUDA_VISIBLE_DEVICES=0,1,2,3 python xxx.py来设置该程序可见的gpu

2. 使用torch.cuda接口

在生成网络对象之前执行：torch.cuda.set_device(0)

3. 使用pytorch的并行GPU接口

net = torch.nn.DataParallel(model, device_ids=[0])

4. 初始化模型时

net = Net.cuda(0)

转载自：PyTorch指定GPU的方法 - 简书 (jianshu.com)