Та же рабочая модель в Keras не улучшается в Pytorch

Question

Я преобразую базовую архитектуру LSTM «многие к одному», чтобы предсказать следующий отдельный элемент в последовательности, записанной в Keras, в Pytorch. Архитектура NN следующая (весь код можно найти здесь ):

model = Sequential()
model.add(LSTM(
    512,
    input_shape=(network_input.shape[1], network_input.shape[2]),
    return_sequences=True
))
model.add(Dropout(0.3))
model.add(LSTM(512, return_sequences=True))
model.add(Dropout(0.3))
model.add(LSTM(512))
model.add(Dense(256))
model.add(Dropout(0.3))
model.add(Dense(n_vocab))
model.add(Activation('softmax'))
model.compile(loss='categorical_crossentropy', optimizer='rmsprop')

Запуск обеих моделей с одинаковыми данными (да, я это явно проверил), обе начинаются со значения потерь ~ 4, но примерно через 100 эпох Keras уже достигла потерь ~ 0,02, что дает желаемые результаты .

Тем не менее, модель Pytorch застряла около ~ 3,4 после 20 эпох. Я перепробовал много вещей:

• Игра с LR: взрывается, когда LR слишком высокий, так что это означает, что по крайней мере параметры обновляются.
• Различные оптимизаторы, SGD, Adam, RMSprop, но одинаковые результаты со всеми.
• Переключение между .view[], .squeeze_ и индексацией при доступе к последнему элементу последовательности.
• Добавление, удаление и изменение нелинейных функций активации и выпадение.
• Удалить ручную инициализацию для x_0 и h_0.

Вот код для моей модели:

class NNP_RNN(nn.Module):
    def __init__(self):
        super(NNP_RNN, self).__init__()
        self.lstm_1 = nn.LSTM(input_size=1, hidden_size=512, batch_first=True)
        self.lstm_2 = nn.LSTM(input_size=512, hidden_size=512, batch_first=True)
        self.lstm_3 = nn.LSTM(input_size=512, hidden_size=512, batch_first=True)

        self.dense_1 = nn.Linear(in_features=512, out_features=256)
        self.dense_2 = nn.Linear(in_features=256, out_features=58)

    def forward(self, x):
        batch_size = x.size(0)
        h_0 = NNP_RNN.init_hidden((1, batch_size, 512))
        c_0 = NNP_RNN.init_hidden((1, batch_size, 512))

        x, _ = self.lstm_1(x, (h_0, c_0))
        x = F.dropout(x, 0.3)

        x, _ = self.lstm_2(x, (h_0, c_0))
        x = F.dropout(x, 0.2)

        _, (x, _) = self.lstm_3(x, (h_0, c_0))
        x = x.squeeze_(0)

        x = self.dense_1(x)
        x = F.dropout(x, 0.1)

        x = self.dense_2(x)

        return x

    @staticmethod
    def init_hidden(dims):
        return torch.zeros(dims, device=device)

И процесс обучения:

optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, verbose=True, patience=5)
criterion = nn.CrossEntropyLoss()

for epoch in range(1, epochs + 1):
    epoch_loss = 0
    epoch_corrects = 0
    for features, labels in tqdm(data, ncols=800):
        features = features.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        batch_size = features.size(0)
        output = model(features)

        loss = criterion(output, labels)
        loss.backward()

        optimizer.step()

        corrects = torch.argmax(output, dim=1)
        corrects = torch.eq(corrects, labels).sum().item()
        epoch_corrects += corrects

        epoch_loss += loss.clone() * batch_size

    epoch_loss /= len(data.dataset)
    epoch_corrects /= len(data.dataset)
    print(f'Loss epoch #{epoch} = {epoch_loss:.10f}, Accuracy = {epoch_corrects}')

    scheduler.step(epoch_loss)