训练 pytorch lstm 模型时笔记本电脑停止运行，而 TensorFlow 模型则正常工作

Question

我在 PyTorch 中有以下 NN 模块

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.emb = nn.Embedding(num_embeddings=10000, embedding_dim=512)
        self.drop1 = nn.Dropout(p=0.25)
        self.lstm = nn.LSTM(input_size=512, hidden_size=32, num_layers=1)
        self.drop2 = nn.Dropout(p=0.25)
        self.dense = nn.Linear(32, 1)
        self.activ = nn.Sigmoid()

    def forward(self, x):
        t1 = self.emb(x)
        t2 = self.drop1(t1)
        outputs, (hidden, cell) = self.lstm(t2)
        t4 = self.drop2(outputs[:,-1,:])
        t5 = self.dense(t4)
        return self.activ(t5)

训练代码如下：

model = Model()
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(3):
    outputs = model(torch.from_numpy(x_train))
    loss = criterion(torch.flatten(outputs).to(torch.float32), torch.flatten(torch.from_numpy(y_train)).to(torch.float32))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

当我大幅降低不同网络层的维度（2 或 4 而不是 512 和 32 等）时，代码工作正常。我这样做是为了调试我的实现并确保它有效。

但是，使用我提供的代码中的给定参数，我的笔记本电脑停止了（鼠标不再移动，没有任何作用，我不得不拔掉笔记本电脑的插头并重新启动它）。在 Google Colab 上执行时同样的情况，出现错误并且会话重置。

我到处都添加了打印，代码似乎停在outputs = model(torch.from_numpy(x_train))处。不过我没有检查前向传递的哪一步。

令人惊讶的是，使用 Tensorflow.Keras 编码的完全相同的模块在我的笔记本电脑和 Google Colab 上都运行良好。

我在这里缺少什么？非常感谢！

我希望培训能够正常进行。

数据下载与处理

import tensorflow as tf
import numpy as np
data = tf.keras.datasets.imdb.load_data(num_words=10000)
train, test = data[0], data[1]
x_train, y_train = train[0], train[1]
x_test, y_test = test[0], test[1]

review_length = 500
from tensorflow.keras.preprocessing import sequence
x_train = sequence.pad_sequences(x_train, maxlen = review_length)
x_test = sequence.pad_sequences(x_test, maxlen = review_length)

Tensorflow Keras 中的相同模型运行良好

from tensorflow.keras.models import Sequential

model = Sequential()
model.add(tf.keras.layers.Embedding(input_dim=10000, output_dim=512, input_length=500))
model.add(tf.keras.layers.Dropout(rate=0.25))
model.add(tf.keras.layers.LSTM(units=32))
model.add(tf.keras.layers.Dropout(rate=0.25))
model.add(tf.keras.layers.Dense(units=1, activation='sigmoid'))

model.compile(optimizer=tf.keras.optimizers.Adam(), loss="binary_crossentropy", metrics=["accuracy"])

model.fit(np.asarray(x_train), y_train, epochs=3, batch_size=256, validation_split=0.2)

Answer 1

问题是在一次调用中将整个训练集传递给模型，简单的解决方案是使用批处理，这就是 Keras 模型起作用的原因。因为

model.fit

自动应用批处理，而在 PyTorch 中必须手动实施批处理。