赞
踩
部分深度学习网络默认是多卡并行训练的,由于某些原因,有时需要指定在某单卡上训练,最近遇到一个,这里总结如下。
上面红框中代码解析:
if torch.cuda.is_available() and ngpu > 1: # 当 torch.cuda.is_available() 为真且 ngpu > 1 时
model = nn.DataParallel(model, device_ids=list(range(ngpu)))
model = nn.DataParallel(model, device_ids=list(range(ngpu))):
此行代码创建了一个 DataParallel包装器,用于在多个GPU上并行处理神经网络模型。DataParallel 是 PyTorch 中的一个模块,它可以将输入数据分割并发送到不同的GPU进行处理,然后汇总结果。
model:要并行化的神经网络模型。
device_ids=list(range(ngpu)):指定要使用的GPU。在这里,它使用了所有可用的GPU,数量上限为指定的 ngpu。
修改好后开始训练,查看显卡使用情况:
以上就是多卡GPU训练与单卡GPU训练相互切换的操作过程,希望能帮到你,谢谢!
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。