【Pytorch】pytorch使用多张GPU进行训练

文章目录

• 使用多张GPU进行训练的代码
• 注意事项

使用多张GPU进行训练的代码

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

model = getModel(args)

if torch.cuda.device_count() > 1:
	print("Let's use", torch.cuda.device_count(), "GPUs!")
	model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])   # device_ids=[0, 1, 2, 3]      
model.to(device)

注意事项

使用多张GPU时，batch_size大于1才能使得多张GPU都工作。