Pytorch 多GPU训练
1 导入库
import torch#深度学习的pytoch平台
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
文章来源:https://www.toymoban.com/news/detail-596253.html
2 指定GPU
2.1 单GPU声明
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
2.2 多GPU声明
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5' #指定GPU编号
device = torch.device("cuda") #创建GPU对象
3 数据放到GPU
x_train = Variable(train,requires_grad=True).to(device=device,dtype=torch.float32) #把训练变量放到GPU
4 把模型网络放到GPU 【重要】
net = DNN(layers)
net = nn.DataParallel(net)
net.to(device=device)
重要:nn.DataParallel
net = nn.DataParallel(net)
net.to(device=device)
- 使用 nn.DataParallel 打包模型
- 然后用 nn.DataParallel 的 model.to(device) 把模型传送到多块GPU中进行运算
torch.nn.DataParallel(DP)
DataParallel(DP)中的参数:
-
module
即表示你定义的模型 -
device_ids
表示你训练时用到的gpu device -
output_device
这个参数表示输出结果的device,默认就是在第一块卡上,因此第一块卡的显存会占用的比其他卡要更多一些。
当调用nn.DataParallel
的时候,input
数据是并行的,但是output loss
却不是这样的,每次都会在output_device
上相加计算
===> 这就造成了第一块GPU的负载远远大于剩余其他的显卡。
DP的优势是实现简单,不涉及多进程,核心在于使用nn.DataParallel
将模型wrap一下,代码其他地方不需要做任何更改。
例子:
5 其他:多GPU并行
加个判断:文章来源地址https://www.toymoban.com/news/detail-596253.html
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model = Model(input_size, output_size) # 实例化模型对象
if torch.cuda.device_count() > 1: # 检查电脑是否有多块GPU
print(f"Let's use {torch.cuda.device_count()} GPUs!")
model = nn.DataParallel(model) # 将模型对象转变为多GPU并行运算的模型
model.to(device) # 把并行的模型移动到GPU上
到了这里,关于Pytorch 多GPU训练的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!