nlp - 翻訳タスクのGPT2トレーニングでバッチサイズを増やす方法は?

Question

機械翻訳タスクに事前トレーニング済みのGPT2モデルを使用するコードを開発しています。データの word-to-id の長さは 91 で、モデル用に次のコードを開発しました。

import torch
from torch.utils.data import DataLoader
from transformers.models.gpt2.modeling_gpt2 import GPT2Model

# data preparation code

def batch_sequences(x, y, env):
    """
    Take as input a list of n sequences (torch.LongTensor vectors) and return
    a tensor of size (slen, n) where slen is the length of the longest
    sentence, and a vector lengths containing the length of each sentence.
    """
    lengths_x = torch.LongTensor([len(s) + 2 for s in x])
    lengths_y = torch.LongTensor([len(s) + 2 for s in y])
    max_length = max(lengths_x.max().item(), lengths_y.max().item())
    sent_x = torch.LongTensor(
        max_length, lengths_x.size(0)).fill_(env.pad_index)
    sent_y = torch.LongTensor(
        max_length, lengths_y.size(0)).fill_(env.pad_index)
    assert lengths_x.min().item() > 2
    assert lengths_y.min().item() > 2

    sent_x[0] = env.eos_index
    for i, s in enumerate(x):
        sent_x[1:lengths_x[i] - 1, i].copy_(s)
        sent_x[lengths_x[i] - 1, i] = env.eos_index

    sent_y[0] = env.eos_index
    for i, s in enumerate(y):
        sent_y[1:lengths_y[i] - 1, i].copy_(s)
        sent_y[lengths_y[i] - 1, i] = env.eos_index

    return sent_x, sent_y, max_length

def collate_fn(elements):
    """
    Collate samples into a batch.
    """
    x, y = zip(*elements)
    x = [torch.LongTensor([env.word2id[w]
                          for w in seq if w in env.word2id]) for seq in x]
    y = [torch.LongTensor([env.word2id[w]
                          for w in seq if w in env.word2id]) for seq in y]
    x, y, length = batch_sequences(x, y, env)
    return (x, length), (y, length), torch.LongTensor(nb_ops)

loader = DataLoader(data, batch_size=1, shuffle=False, collate_fn=collate_fn)
gpt2 = GPT2Model.from_pretrained('gpt2')
in_layer = nn.Embedding(len(env.word2id), 768)
out_layer = nn.Linear(768, len(env.word2id))

parameters = list(gpt2.parameters()) + list(in_layer.parameters()) + list(out_layer.parameters())
optimizer = torch.optim.Adam(parameters)
loss_fn = nn.CrossEntropyLoss()
for layer in (gpt2, in_layer, out_layer):
    layer.train()

accuracies = list()
n_epochs = 5
for i in range(n_epochs):
    for (x, x_len), (y, y_len) in loader:

        x = x.to(device=device)
        y = y.to(device=device)

        embeddings = in_layer(x.reshape(1, -1))
        hidden_state = gpt2(inputs_embeds=embeddings).last_hidden_state[:, :]
        logits = out_layer(hidden_state)[0]
        loss = loss_fn(logits, y.reshape(-1))
        accuracies.append(
            (logits.argmax(dim=-1) == y.reshape(-1)).float().mean().item())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if len(accuracies) % 500 == 0:
            accuracy = sum(accuracies[-50:]) / len(accuracies[-50:])
            print(f'Samples: {len(accuracies)}, Accuracy: {accuracy}')

このコードは、バッチ サイズが 1 の場合にうまく機能しますが、非常に遅いです。バッチ サイズを 1 から 32 に増やしたかったのですが、ディメンションの互換性の問題が発生します。エラーなしでバッチサイズを増やすにはどうすればよいですか?

私のデータは文のペアで構成されています。最初の文は第 1 言語の文で、第 2 文はその第 2 言語への翻訳です。

たとえば、x.shape が (batch_size, 12) であると仮定します (つまり、入力として長さ 12 の 'batch_size' 文があり、y.shape も (batch_size, 12) (翻訳) です)。長さ 90 の to-id ディクショナリで、センテンス内の各単語とそのインデックスに一致します)