Pytorch 多GPU训练

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

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 device
• output_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:# 检查电脑是否有多块GPUprint(f"Let's use {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)# 将模型对象转变为多GPU并行运算的模型

model.to(device)# 把并行的模型移动到GPU上