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
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)
1.使用 nn.dataparallel 打包模型
2.然后用 nn.dataparallel 的 model.to(device) 把模型传送到多块gpu中进行运算
torch.nn.dataparallel(dp)
dataparallel(dp)中的参数:
module即表示你定义的模型device_ids表示你训练时用到的gpu deviceoutput_device这个参数表示输出结果的device,默认就是在第一块卡上,因此第一块卡的显存会占用的比其他卡要更多一些。
当调用nn.dataparallel的时候,input数据是并行的,但是output loss却不是这样的,每次都会在output_device上相加计算
===> 这就造成了第一块gpu的负载远远大于剩余其他的显卡。
dp的优势是实现简单,不涉及多进程,核心在于使用nn.dataparallel将模型wrap一下,代码其他地方不需要做任何更改。
例子:

5 其他:多gpu并行
加个判断:
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上
总结
以上为个人经验,希望能给大家一个参考,也希望大家多多支持代码网。
发表评论