每天讲解一点PyTorch 【18】多卡训练torch.nn.DataParallel

cv.exp

已于 2022-05-22 10:00:12 修改

阅读量1.3k

点赞数

CC 4.0 BY-SA版权

分类专栏： Deep Learning 文章标签： pytorch 人工智能深度学习

于 2022-04-16 19:39:21 首次发布

本文链接：https://blog.csdn.net/forest_world/article/details/124218902

Deep Learning 专栏收录该内容

172 篇文章

订阅专栏

本文介绍了如何在深度学习中利用多个GPU进行模型训练，通过torch.nn.DataParallel实现模型并行，同时展示了如何配置GPU参数，包括设置CUDA_VISIBLE_DEVICES环境变量，并根据需求调整GPU数量和设备ID。此外，还提到了优化器的并行处理以及在代码中处理CPU和GPU的选择。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

使用多GPU训练模型：

net = ......
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

优化器：

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

GPU setting

parser.add_argument(‘–cpu’, action=‘store_true’, help=‘use cpu only’)
parser.add_argument(‘–n_GPUs’, type=int, default=1, help=‘number of GPUs’)
parser.add_argument(‘–gpu_id’, type=int, default=0, help=‘GPU id’)

os.environ[‘CUDA_VISIBLE_DEVICES’]=str(args.gpu_id)
device = torch.device(‘cpu’ if args.cpu else ‘cuda’)

self.feature_loss_module.to(device)

if not args.cpu and args.n_GPUs > 1:
self.feature_loss_module = nn.DataParallel(
self.feature_loss_module, range(args.n_GPUs)
)

![请添加图片描述](https://img-blog.csdnimg.cn/514c2e7481a34b2ca8834a890de65846.png)