 

[mcj]Pytorch单GPU改为多GPU训练

2019-12-23 分类：Pytorch学习评论(0) 赞(320)

原来是：

use_cuda = (args.gpu >= 0) and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
model = model.to(device)

use_cuda = (args.gpu >= 0) and torch.cuda.is_available()

device = torch.device("cuda" if use_cuda else "cpu")

model = model.to(device)

这时用的是单GPU，我们需要改为多GPU，需要在model.to(device)之前加一句

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

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

即可。

赞(320) 打赏

版权声明：本文采用知识共享署名4.0国际许可协议 [BY-NC-SA] 进行授权
文章名称：《[mcj]Pytorch单GPU改为多GPU训练》
文章链接：https://www.machunjie.com/deeplearning/pytorch/626.html
本站资源仅供个人学习交流，请于下载后24小时内删除，不允许用于商业用途，否则法律问题自行承担。

分享到

pytorch