用Pytorch进行多GPU并行训练模型

代码框架为Pytorch，提供一些参考代码给自己和有需要的人

---

前言

主要着重讲dataparallel的使用方法。还有distributed parallel，但我没用过。

---

一、正常的单GPU训练

这个代码很简单，用nvidia-smi查询一下GPU的id，然后输进参数就能用了：

device = torch.device(‘cuda:0’) # for GPU 0
device = torch.device(‘cuda:1’) # for GPU 1
device = torch.device(‘cuda:2’) # for GPU 2

二、多GPU

多GPU分为两种使用情况：数据并行训练和把模型拆开来放不同GPU上训练。

1.data parallel

这个情况就相当于原来可能只是一个batch一个batch（当然每个GPU上的batch数量也可以自己设）的训练，多GPU就可以一次训好几个batch，那计算时间就少了很多，效率基本上和GPU数量呈线性相关。因为这个方法代码好写，用的比较多：

import torch.nn as nn
model = nn.DataParallel(model)

当然，如果你的模型要输入的张量和model不在一张卡上，也会报错，所以对于模型的输入也要进行一些处理，pytorch也贴心地提供了DataLoader和Dataset的数据集的方法，这就要我们把模型的输入像下面写好：

from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, path):
        self.imgs = self.load_imgs(path, debug)

    def load_imgs(self, path, debug):
        # 从 path 加载图像的代码...
        return imgs

    def __getitem__(self, index):
        # 返回一张图像和对应的标签
        # 在这个例子中，我们只返回图像，因为你的数据集没有明确的标签
        return self.imgs[index]

    def __len__(self):
        # 返回数据集中图像的总数
        return len(self.imgs)

# 实例化数据集
train_set = MyDataset(opt.data_path)
test_data = MyDataset(opt.test_path)
train_loader = DataLoader(train_set, batch_size=opt.batch_size, shuffle=True, num_workers = 4)
test_loader = DataLoader(test_data, batch_size=opt.batch_size, shuffle=False, num_workers = 4)

for epoch in range(num_epochs):
    for i, data in enumerate(train_loader):
        # 你的训练代码...

看着还算简单。这里面的shuffle就是问你是否对模型进行打乱，num_workers就是GPU的数量。

2.模型拆在多个GPU上

这个方法就是针对大模型了，我用不到。我就拿pytorch手册的一个例子吧！

import torch
import torch.nn as nn
import torch.optim as optim


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

---

总结

祝看到这篇博文的人都能发顶会。没看到的也能。