PyTorch使用多GPU训练

1. PyTorch使用指定GPU训练

model = Model()
if torch.cuda.is_available():
    model = model.cuda()

上面代码直接使用model.cuda(), PyTorch默认使用从0开始的GPU, 有如下两种方法来指定需要使用的GPU:

• 直接在终端中设定：CUDA_VISIBLE_DEVICES=1 python main.py
• python代码中设定：

  import os
  os.environ["CUDA_VISIBLE_DEVICES"] = "1"
  model = Model()
  if torch.cuda.is_available():
   	model = model.cuda()  #使用第一个GPU
  

2. PyTorch使用指定GPU训练

device_ids = [3, 4, 6, 7]
model = Module()
if torch.cuda.is_available():
    module = torch.nn.DataParallel(model, device_ids=device_ids) # 声明所有可用设备
    model = model.cuda(device=device_ids[0]) # 模型放在主设备

images = images.cuda(device=device_ids[0])   # 训练数据放在主设备
labels = labels.cuda(device=device_ids[0])

使用多GPU训练，model = nn.DataParallel(model)
注意训练/测试过程中 inputs和labels均需加载到GPU中