pytorch多线程加载数据 pytorch多线程训练

一、多卡训练原理

1.为什么要多卡训练：

        简单来说，有两种原因：第一种是模型在一块GPU上放不下，两块或多块GPU上就能运行完整的模型（如早期的AlexNet）。第二种是多块GPU并行计算可以达到加速训练的效果。想要成为“炼丹大师“，多GPU并行训练是不可或缺的技能。

        常见的多卡训练方法有：

        1.模型并行方式（多机多卡）：

        如果模型特别大，GPU显存不够，无法将一个显存放在GPU上，需要把网络的不同模块放在不同GPU上，这样可以训练比较大的网络。

        2.数据并行方式（单机多卡）：

        将整个模型放在一块GPU里，再复制到每一块GPU上，同时进行正向传播和反向误差传播。相当于加大了batch_size。

2.单机多卡训练原理：

        我们主要对数据并行的单机多卡训练进行介绍，在数据并行方式中，我们的训练流程一般如下：

1. 指定主机节点
2. 主机节点划分数据，一个batch数据平均分到每个机器上（数据并行）
3. 模型从主机拷贝到各个机器
4. 每个机器进行前向传播
5. 每个机器计算loss损失
6. 主机收集所有loss结果，进行参数更新
7. 将更新后参数模型拷贝给各个机器

        具体流程图如下：

二、单机多卡训练实现 

1.首先要知道自己的机器有几张显卡

nvidia-smi

2.使用torch.nn.DataParallel(module, device_ids)模块对模型进行包装

        对于torch.nn.DataParallel(module, device_ids)模块，其参数module代表模型，参数device_ids代表并行的GPU id列表。DataParallel的作用是将模型原来的输入(N*M,input_dim) (其中N为显卡数量)拆分成N份，以供N个显卡的GPU进行前向传播，并生成N个大小为(M,output_dim)的输出，最后再将这N个输出在主设备cuda0进行合并，生成完整输出(N*M,output_dim)。

        具体代码示例如下：

device_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # 10张显卡
model = torch.nn.DataParallel(model, device_ids=device_ids) # 指定要用到的设备
model = model.cuda(device=device_ids[0]) # 模型加载到设备cuda0

通过以上原理我们可以看出，DataParallel并没有改变模型的输入和输出。只是将数据分摊给各个GPU进行传播，因此不需要对原来单卡训练的代码进行较大改动。但是，由于最后将N个输出在主设备cuda0进行合并，这也意味着loss计算只由cuda0进行，无法并行计算。

3.解决loss计算负载不均衡问题

        为了解决loss计算负载不均衡的问题，我们就不妨在模型内置loss计算的代码块，最后将所得的loss合并取平均即可。

        具体代码示例如下：

class Net:
    def __init__(self,...):
        # code

    def forward(self, inputs, labels=None)
        # outputs = fct(inputs)
        # loss_fct = ...
        if labels is not None:
            loss = loss_fct(outputs, labels)  # 在训练模型时直接将labels传入模型，在forward过程中计算loss
            return loss
        else:
            return outputs

按照如上代码进行并行运算，在每个GPU上都会计算出一个部分的loss，最后这些loss会被收集到主设备cuda0并合并为长度为N的张量。为了下一步进行backward，我们必须将这个loss张量通过取平均或其他方法转化为一个标量。

 4.将数据指定设备

X_train, y_train = X_train.cuda(device=device_ids[0]), y_train.cuda(device=device_ids[0])

最后，我们通过一个完整的代码来梳理代码逻辑：

import torch
from torchvision import datasets, transforms
import torchvision
from tqdm import tqdm
 
device_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # 可用GPU列表
BATCH_SIZE = 64
 
transform = transforms.Compose([transforms.ToTensor()])
data_train = datasets.MNIST(root = "./data/",
                            transform=transform,
                            train=True,
                            download=True)
data_test = datasets.MNIST(root="./data/",
                           transform=transform,
                           train=False)
 
data_loader_train = torch.utils.data.DataLoader(dataset=data_train,
                                                # 单卡batch size * 卡数
                                                batch_size=BATCH_SIZE * len(device_ids),
                                                shuffle=True,
                                                num_workers=2)
 
data_loader_test = torch.utils.data.DataLoader(dataset=data_test,
                                               batch_size=BATCH_SIZE * len(device_ids),
                                               shuffle=True,
                                               num_workers=2)
 
 
class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = torch.nn.Sequential(
        torch.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(stride=2, kernel_size=2),
    )
        self.dense = torch.nn.Sequential(
            torch.nn.Linear(14 * 14 * 128, 1024),
            torch.nn.ReLU(),
            torch.nn.Dropout(p=0.5),
            torch.nn.Linear(1024, 10)
    )
    def forward(self, x):
        x = self.conv1(x)
        x = x.view(-1, 14 * 14 * 128)
        x = self.dense(x)
        return x
 
 
model = Model()
# 指定要用到的设备
model = torch.nn.DataParallel(model, device_ids=device_ids)
# 模型加载到设备cuda0
model = model.cuda(device=device_ids[0])
 
cost = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
from time import sleep
n_epochs = 50
for epoch in range(n_epochs):
    running_loss = 0.0
    running_correct = 0
    print("Epoch {}/{}".format(epoch, n_epochs))
    print("-"*10)
    for data in tqdm(data_loader_train):
        X_train, y_train = data
        # 指定设备cuda0
        X_train, y_train = X_train.cuda(device=device_ids[0]), y_train.cuda(device=device_ids[0])
        outputs = model(X_train)
        _,pred = torch.max(outputs.data, 1)
        optimizer.zero_grad()
        loss = cost(outputs, y_train)
 
        loss.backward()
        optimizer.step()
        running_loss += loss.data.item()
        running_correct += torch.sum(pred == y_train.data)
    testing_correct = 0
    for data in data_loader_test:
        X_test, y_test = data
        # 指定设备1
        X_test, y_test = X_test.cuda(device=device_ids[0]), y_test.cuda(device=device_ids[0])
        outputs = model(X_test)
        _, pred = torch.max(outputs.data, 1)
        testing_correct += torch.sum(pred == y_test.data)
    print("Loss is:{:.4f}, Train Accuracy is:{:.4f}%, Test Accuracy is:{:.4f}".format(torch.true_divide(running_loss, len(data_train)),
                                                                                      torch.true_divide(100*running_correct, len(data_train)),
                                                                                      torch.true_divide(100*testing_correct, len(data_test))))
torch.save(model.state_dict(), "model_parameter.pkl")

三、多卡训练的模型保存

      多卡训练中的模型的保存与加载，与单GPU的方式有所不同。这里通通将参数以cpu的方式保存, 因为如果是保存的GPU参数，pth文件中会记录参数所属的GPU号，加载时会加载到相应的GPU上，这样就会导致如果GPU数目不够时会在加载模型时报错。

        保存代码示例如下：

torch.save(model.module.cpu().state_dict(), "model.pth")