Формы mat1 и mat2 не могут быть перемножены (100×200 и 100×9922)

Я пытаюсь создать языковую модель BiLSTM и испытываю некоторые проблемы.

Модель

class BiLSTM(nn.Module):
  def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers, dropout_rate, tie_weights):
    super().__init__()

    self.num_layers = num_layers
    self.hidden_dim = hidden_dim
    self.embedding_dim = embedding_dim

    self.embedding = nn.Embedding(vocab_size, embedding_dim)
    self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=num_layers, 
                        dropout=dropout_rate, batch_first=True, bidirectional=True)
    self.dropout = nn.Dropout(dropout_rate)
    self.linear = nn.Linear(hidden_dim*2, vocab_size)

    if tie_weights:
      # Входной и скрытый слой должны быть одного размера для связывания весов
      assert embedding_dim == hidden_dim, 'Невозможно связать веса, проверьте размеры'
      self.linear.weight = self.embedding.weight

    self.init_weights()

    self.hidden = None
    self.cell = None

  def forward(self, x, hidden):
    hidden_0 = torch.zeros(2*num_layers, batch_size, self.hidden_dim).to(device)
    cell_0 = torch.zeros(2*num_layers, batch_size, self.hidden_dim).to(device)
    output  = self.embedding(x)
    output, (h, c) = self.lstm(output, (hidden_0, cell_0))
    output = self.dropout(output)
    output = self.linear(output)
    return output, (h, c)

  def init_weights(self):
    init_range_emb = 0.1
    init_range_other = 1/math.sqrt(self.hidden_dim)
    self.embedding.weight.data.uniform_(-init_range_emb, init_range_emb)
    self.linear.weight.data.uniform_(-init_range_other, init_range_other)
    self.linear.bias.data.zero_()

  def init_hidden(self, batch_size):
        hidden = torch.zeros(self.num_layers, batch_size, self.hidden_dim)
        cell = torch.zeros(self.num_layers, batch_size, self.hidden_dim)
        return hidden, cell

  def detach_hidden(self, hidden):
        hidden, cell = hidden
        hidden = hidden.detach()
        cell = cell.detach()
        return hidden, cell

Параметры модели

vocab_size = len(vocab)
embedding_dim = 100
hidden_dim = 100
num_layers = 5
dropout_rate = 0.4
tie_weights = True
model = BiLSTM(vocab_size, embedding_dim, hidden_dim, num_layers, dropout_rate, tie_weights)
model.to(device)

Обучение:

import copy
import time

criterion = nn.CrossEntropyLoss()
lr = 20.0  # скорость обучения
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

def train(model: nn.Module) -> None:
    model.train()  # включить режим обучения
    total_loss = 0.
    log_interval = 200
    start_time = time.time()

    hidden = model.init_hidden(batch_size)

    num_batches = len(train_data) // bptt
    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):
        hidden = model.detach_hidden(hidden)
        data, targets = get_batch(train_data, i)
        seq_len = data.size(0)
        output, hidden = model(data, hidden)
        loss = criterion(output.view(-1, vocab_size), targets)

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()

        total_loss += loss.item()

        if batch % log_interval == 0 and batch > 0:
            lr = scheduler.get_last_lr()[0]
            ms_per_batch = (time.time() - start_time) * 1000 / log_interval
            cur_loss = total_loss / log_interval
            ppl = math.exp(cur_loss)
            print(f'| эпоха {epoch:3d} | {batch:5d}/{num_batches:5d} батчей | '
                  f'lr {lr:02.2f} | мс/батч {ms_per_batch:5.2f} | '
                  f'потери {cur_loss:5.2f} | ppl {ppl:8.2f}')
            total_loss = 0
            start_time = time.time()

Оценка:

def evaluate(model: nn.Module, eval_data: Tensor) -> float:
    model.eval()  # включить режим оценки
    total_loss = 0.
    with torch.no_grad():
        for i in range(0, eval_data.size(0) - 1, bptt):
            data, targets = get_batch(eval_data, i)
            seq_len = data.size(0)
            output = model(data)
            output_flat = output.view(-1, vocab_size)
            total_loss += seq_len * criterion(output_flat, targets).item()
    return total_loss / (len(eval_data) - 1)

Блок обучения:

best_val_loss = float('inf')
epochs = 50
best_model = None

for epoch in range(1, epochs + 1):
    epoch_start_time = time.time()
    train(model)
    val_loss = evaluate(model, val_data)
    val_ppl = math.exp(val_loss)
    elapsed = time.time() - epoch_start_time
    print('-' * 89)
    print(f'| конец эпохи {epoch:3d} | время: {elapsed:5.2f}с | '
          f'потери на валидации {val_loss:5.2f} | ppl на валидации {val_ppl:8.2f}')
    print('-' * 89)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_model = copy.deepcopy(model)

    scheduler.step()

Моя проблема в том, что я получаю следующую ошибку.

RuntimeError                              Traceback (most recent call last)
<ipython-input-345-453c3f2a9cad> in <cell line: 5>()
      5 for epoch in range(1, epochs + 1):
      6     epoch_start_time = time.time()
----> 7     train(model)
      8     val_loss = evaluate(model, val_data)
      9     val_ppl = math.exp(val_loss)

4 frames
<ipython-input-343-16d7ac1074e2> in train(model)
     20         data, targets = get_batch(train_data, i)
     21         seq_len = data.size(0)
---> 22         output, hidden = model(data, hidden)
     23         loss = criterion(output.view(-1, vocab_size), targets)
     24 

/usr/local/lib/python3.9/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)
   1499                 or _global_backward_pre_hooks or _global_backward_hooks):
   1500
   1501     return forward_call(*args, **kwargs)
   1502                 or _global_backward_pre_hooks или _global_backward_hooks):
   1503 
   1504 # Вызывайте только в том случае, если jit не используется
   1505 full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.9/dist-packages/torch/nn/modules/linear.py in forward(self, input)
    112 
    113     def forward(self, input: Tensor) -> Tensor:
----> 114         return F.linear(input, self.weight, self.bias)
    115 
    116     def extra_repr(self) -> str:

RuntimeError: размеры mat1 и mat2 не могут быть перемножены (100x200 и 100x9922)

Проблема, похоже, связана с размерами модели, но я не знаю, как настроить модель, чтобы решить эту проблему.

Любая помощь будет полезной, спасибо.

ИЗМЕНЕНИЕ

Функция get_batch возвращает

(tensor([[9893,    0,  171,    1,    1,  180,    4,  795,   13,   26],
        [9894, 1192,   58,    0, 9325,  190,  140,    4,    2,    0],
        [9896,    2,  109, 5437,   20,   12, 1022,    5,    6,    4],
        [9897,    2,   52, 1446,    5,  474,    9,   38,  192,    5],
        [9898,    3, 2850,   30, 3489,  655,   10,  212, 3915,  300],
        [9902,    1, 2150,    1, 1791, 1534,  141,   12,    4, 2853],
        [9903, 2526,    9,  161,   40,   17,    7,    5,  339,    4],
        [9904,   42,   10,   14,  775, 1931,  151, 1786, 1281,  893],
        [9905,   32, 1574,    2, 2088,  569,  154,   19,  799, 3353],
        [9906,  222,   43,    9,  342,  499,  290,    1,    0, 9253]],
       device="cuda:0"), tensor([9894, 1192,   58,    0, 9325,  190,  140,    4,    2,    0, 9896,    2,
         109, 5437,   20,   12, 1022,    5,    6,    4, 9897,    2,   52, 1446,
           5,  474,    9,   38,  192,    5, 9898,    3, 2850,   30, 3489,  655,
          10,  212, 3915,  300, 9902,    1, 2150,    1, 1791, 1534,  141,   12,
           4, 2853, 9903, 2526,    9,  161,   40,   17,    7,    5,  339,    4,
        9904,   42,   10,   14,  775, 1931,  151, 1786, 1281,  893, 9905,   32,
        1574,    2, 2088,  569,  154,   19,  799, 3353, 9906,  222,   43,    9,
         342,  499,  290,    1,    0, 9253, 9908,    6,  388,   31,   28,    0,
         559,  835,    7,   30], device="cuda:0"))

На первый взгляд, мне интересно, не является ли это чем-то простым: согласно документации по nn.LSTM, batch_first – если True, то входные и выходные тензоры предоставляются в виде (batch, seq, feature) вместо (seq, batch, feature). Обратите внимание, что это не относится к скрытым или ячеистым состояниям. См. разделы Входные/Выходные данные ниже для получения дополнительных сведений. По умолчанию: False. Если вы установите это значение на False, то размеры вашего батча должны быть 200×100, и это должно сработать (возможно?).