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有关。