pytorch多GPU数据并行模式 踩坑指南
转载声明:
转自https://blog.csdn.net/yuuyuhaksho/article/details/87560640
仅备份用作自己学习使用。
————————————————
版权声明:本文为CSDN博主「Edward Tivrusky IV」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/yuuyuhaksho/article/details/87560640
正文
之前用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也很重要。
- 在init函数中预定义的独立的神经网络层,在forward函数里形成神经网络结构
- 要保证这个函数里的tensors正确地分成了需要计算梯度的和不需要计算梯度的
- 这个函数的输入和输出变量对于需要平行计算的模型至关重要!!
踩坑点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