pytorch多GPU训练方式 Multi GPU

1、首先定义模型

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        pass
    def forward(self, x):
        pass

2、利用DataParallel分布式存储数据和模型

gpuID = '0,1,2,3,4,5,6,7'
device_ids = [int(id) for id in gpuID.split(',')]

import os
os.environ['CUDA_VISIBLE_DEVICES'] = gpuID

device = torch.device('cuda: 0')

import torch
model = Net()
model = torch.nn.DataParallel(model, device_ids=device_ids)
model.to(device)
criterion.to(device)

input, output = input.to(device), output.to(device)

output_ = model(input)

loss = criterion(output, output_)

3、查看nvidia-smi命令即可发现模型训练已经自动分配在全部显卡上

4、貌似指定偶数块GPU没问题，实验过程中，如果指定奇数块GPU，如4，5，6，发现模型只占用了4，5号GPU；而指定4，5，6，7时，模型会占用4，5，6，7号GPU。没搞懂原因。搞懂了，跟batchsize有关。