使用DataParallel,batchsize需要大于GPU数量
单GPU代码
device = torch.device("cuda")
model.to(device)
改为多张GPU代码
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1' #代码最前面加入
device = torch.device('cuda')
model = toech.nn.DataParallel(model, [0,1])
model.to(device)
单GPU代码
device = torch.device("cuda")
model.to(device)
改为多张GPU代码
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1' #代码最前面加入
device = torch.device('cuda')
model = toech.nn.DataParallel(model, [0,1])
model.to(device)