单机多卡训练

weixin_45833431

已于 2023-06-10 21:03:26 修改

阅读量105

点赞数

CC 4.0 BY-SA版权

文章标签： pytorch 深度学习神经网络

于 2023-04-18 22:08:34 首次发布

本文链接：https://blog.csdn.net/weixin_45833431/article/details/130232534

该段代码设置了使用GPU4,5,6,7进行多卡测试，主显卡设为4。定义了DCENet模型，并在显卡数量大于1时使用DataParallel进行数据并行处理，确保模型能在多个GPU上自动分配运行。

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

#---多卡测试---在这里设置显卡号
gpu_ids = [4, 5, 6, 7]#设置用哪几个显卡
device = torch.device("cuda:4")#设置哪个显卡为主显卡
model =DCENet(n=8, return_results=[4, 6, 8])
if len(gpu_ids) > 1:
    model = torch.nn.parallel.DataParallel(model , device_ids = gpu_ids)#重点：自动分配数据
model.to(device)
#---多卡测试