本文介绍了如何在PyTorch中使用nn.DataParallel封装模型并将模型部署到多个GPU上，同时讨论了torch.nn.DataParallel与torch.nn.parallel.DistributedDataParallel之间的区别，尤其是在处理多GPU和batchsize时的注意事项。
如果有多个GPUs，可以使用nn.DataParallel封装模型，然后使用model.to(device)把模型放在GPUs上。
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
model.to(device)
如果有3个GPUs，batch size为30，则每个gpu batch 10个输入。
挖个坑：
看一下torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel的区别
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
