PyTorch中指定GPU块号的方法

1、 改变系统变量使得仅目标显卡可见

1. export CUDA_VISIBLE_DEVICES=0,1,2...

2. os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2...'

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

   这里的“0,1,2”指的是gpu的编号，一般情况下是从0开始。

2、使用torch.cuda接口

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

3、使用pytorch的并行GPU接口

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

4、初始化模型时

net = Net.cuda(0)