pytorch多GPU数据并行模式 踩坑指南

之前用pytorch尝试写了个文本生成对抗模型seqGAN，相关博文在这里。

在部署的时候惊喜地发现有多块GPU可供训练用，于是很天真地决定把之前写的单GPU版本改写成DataParallel的方式（内心os：介有嘛呀）。于是开始了从入门到（几乎）放弃的踩坑之路。

为了和大家共同进步，我把自己的经验分享一下，欢迎一起来踩坑。

首先说明，我用的pytorch版本虽然不是嘎嘣新的1.0，但是是稳定版本0.4，而且这期间调研的结果，1.0版本并没有解决数据并行中所有大家遇到的问题。如果有童鞋用的1.0，可以参考这份指南。

另外说明，这份指南适合已经有一定pytorch经验的童鞋。指南的重点是介绍自定义神经网络在DataParallel模式下的工作方式，另外涉及少量自定义RNN类网络的结构。

torch tensor类型

当使用单device环境时，对tensor类型的设置虽然重要但并不致命。因为不涉及tensor的拆分，所以基本只要使用各种pytorch默认设置就行了。但是当涉及数据并行，就需要更了解模型不同位置需要的不同tensor类型。

pytorch的tensor有两个大类：cpu和cuda tensor。每个类都对应了不同的数据类型，具体列表在这里和这里。

在生成tensor时，可以选择通过定义device直接在指定的设备上生成：

tensor = torch.tensor([2, 4], dtype=torch.float64, device=torch.device('cuda'))

这段代码会在GPU上生成一个值为[2.,4.]，大小为 1行 x 2列，数据类型为64位浮点型的tensor。

也可以先在CPU上生成，再放到GPU上:

tensor = torch.tensor([2, 4], dtype=torch.float64).to(torch.device('cuda'))

或者：

tensor = torch.tensor([2, 4], dtype=torch.float64).cuda()

也可以直接指定生成类型(cpu或者cuda tensor)：

tensor = torch.DoubleTensor([2, 4])

会直接生成一个cpu类型的tensor，

tensor = torch.cuda.DoubleTensor([2, 4])

会直接生成一个cuda类型的tensor。

转换tensor的类型(cpu/cuda，数据类型)，还可以用：

tensor = torch.cuda.DoubleTensor([2, 4]).type(torch.LongTensor)

torch.tensor接受已经存在的数据（在我们的例子中是[2,4]这个列表），并且会生成一份copy。如果需要从比如numpy类型的数组直接生成torch.tensor，而不需要copy，可以用：

tensor = torch.as_tensor(np.array([2,4]))

踩坑点1：注意大写Tensor和小写tensor的区别。

torch.Tensor也可以接收已经存在的数据，比如

tensor = torch.tensor([2, 4]) 会生成一个torch.int64类型的tensor，而

tensor = torch.Tensor([2, 4]) 也会生成一个值为[2.,4.]，大小为 1行 x 2列的tensor，但是类型是torch.float32。

tensor = torch.Tensor(2, 4) 会生成一个大小为 2行 x 4列的空白矩阵。

在最近几版pytorch中，Tensor的功能被拆成了tensor和empty两种，前者接收已有的数据，后者接收数据维度信息生成新矩阵。

踩坑点2：用empty生成新矩阵时，矩阵的值并不是默认为0。

torch.empty(2,4).sum() 会返回随机值。

生成零矩阵的正确方式是：

tensor = torch.empty(2, 4).zero_()

或者

tensor = torch.zero_(torch.empty(2, 4))

踩坑点3：注意 torch.set_default_tensor_type的设置。

tensor = torch.tensor([2, 4]) 即使在cuda环境也默认返回 cpu 类型。

但当设置 torch.set_default_tensor_type(torch.cuda.FloatTensor) 后，同样的代码会返回cuda类型的tensor.

pytorch中自定义神经网络的代码结构

首先创建一个继承nn.Module的子类。

在这个类中__init__()函数是必要的，除了执行super的init函数以外，一般在这里预定义需要的神经网络层。另外网上有很多代码示例在这个函数里还放了包括hidden层和预测结果在内的很多tensor/变量。我的建议是：请仔细考虑你需要把什么tensor绑定在一个特定的实例里，而需要什么tensor具有scalability（允许平行计算）。这个踩坑点后面我还会提到。

除了__init__()函数以外，自定义神经网络类中最重要的是forward()函数。这个函数在前馈过程(forward pass)中被隐性调用（调用时只是指明了类的名称，并没有显性调用forward()函数。见下面的代码示例中的train()部分），而且对于BP也很重要。

1. 在init函数中预定义的独立的神经网络层，在forward函数里形成神经网络结构
2. 要保证这个函数里的tensors正确地分成了需要计算梯度的和不需要计算梯度的
3. 这个函数的输入和输出变量对于需要平行计算的模型至关重要！！

踩坑点4：batch first 还是 batch second：在tensor输入维度上可以选择第一位是batch size，或者第二位是batch size。理论上说这是个人习惯问题，只要前后统一就可以；但是在pytorch中内置的lstm层只接受batch second 的hidden层tensor。虽然在lstm层或padding层上可以定义batch_first=True，但是这只定义了输入tensor；hidden层仍然必须是batch second 格式。具体参考这里。

所以请仔细考虑lstm层的输入/输出/hidden/cell的数据格式。如下面代码示例：所有tensor（包括hidden）都采用了batch_first=True的格式，但hidden tuple在输入lstm层以前被转置了。

踩坑点5：当模型是nn.DataParallel类型时，在执行model.forward()函数时同一batch会被分配到不同GPU上平行计算。拆分的维度默认为第一维(dim=0)，但可以设置为其他维度进行拆分（比如如果你习惯所有的tensor都用batch second 的格式，就可以设置拆分维度为dim=1）。前提是所有输入tensor都必须是cuda类型。cpu类型的输入只会被原样拷贝到每个实例中而不会被拆分。

如果选择DataParallel方式，所有与batch size 有关的数据（比如示例代码中的sentence_lengths 和 hidden）都必须和input tensor （示例代码中是sentence）一起作为forward()函数的入参，而不能作为绑定在实例上的比如self.hidden形式存在。如果与实例绑定，根据GPU个数的数据拆分会引起错误。

以示例代码为例：如果在train()函数中生成的hidden格式为(batch size, 1, hidden size) = （128, 1, 48)，GPU个数为2个，DataParallel的拆分维度为dim=0：会在device:0上生成实例并拷贝到device:1上，而调用forward()函数时，两个实例会分别获取格式为(64, 1, 48)的hidden作为输入。

同样道理，所有forward()函数的输出结果都会自动在拆分维度上被concatenate回完整的输出格式。所以如果有跟batch_size有关的输出，比如预测结果，也必须以forward()函数的输出的格式返回，而不能绑定在实例上。

所以请仔细考虑DataParrallel模型的forward()函数的输入和输出变量，以及self.*格式的变量。

另外请注意batch的大小。当batch小于GPU个数时，系统会报错。

踩坑点6：注意多个DataParrallel类型的引用和嵌套关系。

如果新的自定义神经网络是基于已有神经网络class建立的（比如A网络包裹了B网络，在B的基础上又增加了新的网络层），需要注意是否有重复定义nn.DataParallel类型的情况。

pytorch规定：在使用nn.DataParallel类型时，首个实例必须建立在device:0上。当出现嵌套的class重复定义了nn.DataParallel时，第一次拆分外层class会把实例拷贝并抛到所有GPU上执行（比如device:1)，而device:1上运行的部分会生成内层class，引起 RuntimeError: all tensors must be on devices[0]。细节参考这里。

关于padding层

对于RNN来说，由于经常遇到序列长度相差很多的场景（比如句子长度）。通过padding可以大幅提高训练效率。在pytorch中，padding是通过pack_padded_sequence 和 pad_packed_sequence 实现的，前者是padding，后者是还原。

基础使用教程参考这里。

踩坑点7：输入文本长度的排序：pack_padded_sequence 只接受lengths为降序排列的batch。讨论在这里。

如果文本量很大，可以考虑在forward()函数里进行排序。

踩坑点8：在数据并行模式下，输出文本长度total_length需要显性定义，不能用默认值。

padding还原层 pad_packed_sequence 默认按照输入的最长序列还原。在单device环境下，还原层能够正确根据最长序列还原padding。但是在多GPU环境下，每个GPU上的还原层只能看到当前device的最长序列，造成在concatenate所有实例的forward()结果时报错。

这个问题在今年4月的一个feature request中得到了解决。在数据并行下使用还原层时，需要定义total_length参数，统一每个device上的序列长度。

踩坑点9：序列长度的数据格式必须为int64类型的cpu tensor。

说实话这是一个很诡异的设定。pytorch开发者给出的答案是:

apaszke commented on May 10:
It’s a feature. Lengths are used in various conditionals and storing them on CUDA would impose unnecessary overheads.

在单device环境下没有问题，可以对forward()函数输入一个python list 类型的lengths 参数，padding层会自动转化为cpu类型tensor。

但是在多GPU环境下就比较尴尬了。lengths与batch size 相关，需要和input tensor一起被拆分，所以在调用model.forward()函数之前必须是一个cuda类型的tensor (参见踩坑点5）。然后在forward()函数内部需要被转换到cpu类型才能输入给padding层。

自定义RNN 的示例代码

import torch
import torch.nn as nn

DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

class LSTMTest(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(10, 32)
        self.lstm = nn.LSTM(32, 5, batch_first=True)
        self.hidden2tag = nn.Linear(5, 10)
        self.logSoftmax = nn.LogSoftmax(dim=2)

    def init_hidden(self, batch_size=1):
        return (torch.empty(batch_size, 1, 5, device=DEVICE).normal_(),
                torch.empty(batch_size, 1, 5, device=DEVICE).normal_())

    def forward(self, sentence, sentence_lengths, hidden):
        sentence_lengths = sentence_lengths.type(torch.LongTensor)
        embeds = self.embedding(sentence.long())
        embeds = torch.nn.utils.rnn.pack_padded_sequence(embeds, 
				 sentence_lengths.to(torch.device('cpu')), batch_first=True)
        hidden0 = [x.permute(1,0,2).contiguous() for x in hidden]
        lstm_out, hidden0 = self.lstm(embeds, hidden0)
        lstm_out, _ = torch.nn.utils.rnn.pad_packed_sequence(lstm_out, 
					  batch_first=True, total_length=sentence.shape[1])
        tag_space = self.hidden2tag(lstm_out)
        tag_scores = self.logSoftmax(tag_space)
        return tag_scores, tag_space

def train():
    try:
        print('number of GPUs available:{}'.format(torch.cuda.device_count()))
        print('device name:{}'.format(torch.cuda.get_device_name(0)))
    except:
        pass
    sentence = torch.rand(100, 8, device=DEVICE)
    sentence = torch.abs(sentence * (10)).int()
    sentence_lengths = [sentence.shape[1]] * len(sentence)

    model = LSTMTest()
    model = nn.DataParallel(model)
    model.to(DEVICE)
    params = list(filter(lambda p: p.requires_grad, model.parameters()))
    criterion = nn.NLLLoss()
    optimizer = torch.optim.SGD(params, lr=0.01)
    batch_size = 6
    for epoch in range(3):
        pointer = 0
        while pointer + batch_size <= len(sentence):
            x_batch = sentence[pointer:pointer+batch_size]
            x_length = torch.tensor(sentence_lengths[pointer:pointer+batch_size]).to(DEVICE)
            y = x_batch
            hidden = model.module.init_hidden(batch_size=batch_size)
            y_pred, tag_space = model(x_batch, x_length, hidden)
            loss = criterion(y_pred.view(-1,y_pred.shape[-1]), y.long().view(-1))
            optimizer.zero_grad()
            loss.backward(retain_graph=True)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.25)
            optimizer.step()
            pointer = pointer + batch_size