torch如何指定多GPU训练

训练的时候可以使用 .to(device)和.cuda()，注意两者的区别。

一、.to(device)，可以指定使用CPU或GPU

#使用单GPU或者CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#也可以手动指定使用GPU或CPU，如下:
# device=torch.device("CPU")
# device=torch.device("GPU")
model.to(device)#将模型加载到相应的设备中

#使用多GPU
if torch.cuda.device_count() > 1:
  model = nn.DataParallel(model，device_ids=[0,1,2])
model.to(device)
#【注】：这里必须使用0块显卡，不然会报错

#如果使用所有的GPU，那么：
if torch.cuda.device_count() > 1:
  model = nn.DataParallel(model)
model.to(device)

二、cuda()，只能指定GPU

#指定单个GPU
import os
os.environ['CUDA_VISIBLE_DEVICES']='2'
model.cuda()

#指定多个GPU
from torch.nn import DataParallel
#上面这一行一定要加，因为没加一直报错都开始怀疑torch版本是不是有问题了哈哈
import os
os.environ['CUDA_VISIBLE_DEVICES']='0,2,3,4'#代表使用的第一块，第三块，第四块和第五块卡
device_ids = [0,2,3,4]
model  = torch.nn.Dataparallel(model, device_ids =device_ids)
model = model.cuda()

#如果默认使用所有的GPU，那么：
model = torch.nn.Dataparallel(model )

#也可以在终端指定GPU
CUDA_VISIBLE_DEVICES=0,2,3,4 python ***.py

三、分布式训练torch.distributed.launch

-m torch.distributed.launch --nproc_per_node=NGPUS --master_port=$PORT
#NGPUS代表使用了几块GPU卡
#上面那条语句是按顺序选择卡，如果想要指定固定的卡，可以在跑程序时先在终端指定卡号然后再输入上面的语句：
#终端指定卡号：
export CUDA_VISIBLE_DEVICES=0,2