训练PyTorch LSTM模型时设备死机，TensorFlow版本运行正常求助

核心问题 #

你的PyTorch代码直接用全量训练数据（25000条样本）一次性喂给模型，导致内存占用远超硬件承载上限，而TensorFlow Keras的model.fit默认按指定batch_size分批处理数据，所以能正常运行。以下是具体修复点：

1. 实现批量训练逻辑 #

这是最关键的修复，用PyTorch的DataLoader分批加载数据，和Keras的batch_size=256保持一致：

from torch.utils.data import TensorDataset, DataLoader

# 提前将数据转换为PyTorch Tensor并构建数据集
train_dataset = TensorDataset(
    torch.from_numpy(x_train).long(),  # Embedding层要求输入为长整型
    torch.from_numpy(y_train).float()
)
# 按批次加载，shuffle=True保证训练时样本打乱
train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True)

# 修改训练循环为批量迭代
model = Model()
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(3):
    total_loss = 0.0
    for x_batch, y_batch in train_loader:
        # 前向传播
        outputs = model(x_batch)
        # 计算损失，无需重复flatten和类型转换
        loss = criterion(outputs.squeeze(), y_batch)
        
        # 反向传播与优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}")

2. 优化数据类型与维度匹配 #

• 提前转换数据类型，避免在训练循环中重复执行to()操作，减少内存开销
• 用outputs.squeeze()去掉冗余维度，让输出和标签维度对齐

3. 启用GPU加速（可选但关键） #

Keras会自动检测并使用GPU，PyTorch需要手动指定设备，大幅降低训练时间：

# 检测可用设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Model().to(device)

# 训练时将批次数据移到对应设备
for epoch in range(3):
    total_loss = 0.0
    for x_batch, y_batch in train_loader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)
        
        outputs = model(x_batch)
        loss = criterion(outputs.squeeze(), y_batch)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    print(f"Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}")

4. LSTM输入维度兼容（可选） #

PyTorch的LSTM默认输入格式为(seq_len, batch_size, input_size)，而Keras默认是(batch_size, seq_len, input_size)，你的代码当前逻辑可以正常运行，但如果遇到维度报错，可在forward方法中转置输入：

def forward(self, x):
    t1 = self.emb(x)  # 原始shape: (batch_size, seq_len, 512)
    t1 = t1.transpose(0, 1)  # 转置为(seq_len, batch_size, 512)
    t2 = self.drop1(t1)
    outputs, (hidden, cell) = self.lstm(t2)
    t4 = self.drop2(outputs[-1, :, :])  # 取最后时间步输出
    t5 = self.dense(t4)
    return self.activ(t5)

内容的提问来源于stack exchange，提问作者Brahim Khalil Abid