pytorch使用DataParallel加速（包括RNN疑难杂症处理）

DataParallel的基本使用方法很简单，只需设置device_ids即可，如下所示：

device_ids = [0, 1, 2, 3]
model = torch.nn.DataParallel(model, device_ids=device_ids)

device_ids为你要使用的GPU号。如果你未使用DataParallel之前用的便是单GPU进行训练，那么对于数据不需要额外的操作，否则，你需要将模型的输入数据转移到cuda上，如：

# 此处device与device_ids无关，你可以设置device = torch.device("cuda:0")
input = input.to(device)

如果顺利的话，简单的两步就可以实现加速了。

然而，由墨菲定律可得：凡是可能出错的事就一定会出错。常见问题如下。

问题1：如果model里定义了一个函数，如初始化函数init_hidden等，并已实现DataParallel，在train函数里该如何调用？

class Model(nn.Module):
    def __init__(self, ):
        pass

    def forward(self, ):
        pass

    def init_hidden(self, ):
        pass


model = Model()
model = nn.DataParallel(model, device_ids=device_ids)

#调用init_hidden函数,需加上.module
model.module.init_hidden()

解释：

nn.DataParallel是一个Pytorch的nn.Module，模型和优化器都需要通过使用.module来得到实际的模型和优化器。​​​​​​​

问题2： 使用DataParallel来加速RNN、LSTM、GRU等模型，遇到诸如RuntimeError: Expected `len(lengths)` to be equal to batch_size, but got 128 (batch_size=43)问题。

首先RNN中使用pack_padded_sequence和pad_packed_sequence时，需注意以下几点：

1、pack_padded_sequence只接受lengths为降序排列的batch，可通过DataLoader函数的collate_fn参数来进行排序。

2、在数据并行模式下，输出文本长度total_length需要显性定义。

#设置total_length
gru_out, hidden = self.gru(embs, hidden0)
gru_out, lengths = pad_packed_sequence(gru_out, batch_first=True, total_length=seq.shape[1])

3、在执行model.forward()函数时同一batch会被分配到不同的GPU上进行计算（只有输入为cuda上的tensor才会被拆分，cpu类型的输入只会被原样拷贝到每个实例中而不会被拆分）。

因此，如果选择DataParallel方式，所有与batch_size有关的数据都必须和input tensor一起作为forward()函数的参数，而不能以绑定在实例上的形式存在，如self.hidden。如果在train函数中生成的hidden格式为(batch size, num_layers, hidden size) = （128, 2, 100)，GPU个数为2个，DataParallel的拆分维度为dim=0：会在device:0上生成实例并拷贝到device:1上，而调用forward()函数时，两个实例会分别获取格式为(64, 2, 100)的hidden作为输入。

4、序列长度的数据格式必须为int64类型的cpu tensor。

在单卡环境下，对forward()上输入list类型的lengths参数，padding层会自动转化为cpu类型tensor。多卡时，lengths与batch_size有关，需要和input tensor一起被拆分，所以在调用model.forward()函数之前必须是一个cuda类型的tensor, 然后在forward()函数内部需要被转换到cpu类型才能输入给padding层。

示例代码：

class Model(nn.Module):

    def __init__(self, n_items, hidden_size, embedding_dim, n_layers=1):
        super(Model, self).__init__()
        self.n_items = n_items
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.embedding_dim = embedding_dim
        self.emb = nn.Embedding(self.n_items + 1, self.embedding_dim, padding_idx=0)
        self.gru = nn.GRU(self.embedding_dim, self.hidden_size, self.n_layers, batch_first=True)
        self.device = device

    def forward(self, seq, lengths, hidden):
        # 可以输出hidden的shape和divice，发现数据在多个gpu上，且每个gpu的batchsize为总batchsize除以卡的个数
        # print(hidden.shape)
        # print(lengths.device)
        embs = self.emb(seq)
        embs = pack_padded_sequence(embs, lengths.to(torch.device("cpu")), batch_first=True)
        hidden0 = hidden.permute(1, 0, 2).contiguous()
        gru_out, hidden = self.gru(embs, hidden0)
        gru_out, lengths = pad_packed_sequence(gru_out, batch_first=True, total_length=seq.shape[1])
        ht = hidden[-1]
        return ht

    def init_hidden(self, batch_size):
        return torch.zeros((batch_size, self.n_layers, self.hidden_size), requires_grad=True).to(self.device)


def train():
    model = Model(item_nuniq, hidden_size, emb_dim, n_layers).to(device)
    model = nn.DataParallel(model, device_ids=device_ids)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    for epoch in range(epochs):
        model.train()
        for i, (user, hist_click, target, lens) in enumerate(train_loader):
            lens = torch.tensor(lens).to(device)
            hist_click, target = hist_click.to(device), target.to(device)
            hidden = model.module.init_hidden(batch_size=128)
            output = model(hist_click, lens, hidden)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    

附：

 参考链接：

知乎 Pytorch单机并行训练--会飞的闲鱼

知乎 Pytorch的nn.DataParallel -- Mario

知乎 pytorch多gpu并行训练 --link-web

尤其指出：CSDN pytorch多GPU数据并行模式 踩坑指南，本文多参考自该文