PyTorch时序预测LSTM模型性能劣于TensorFlow，求差异排查建议

我有一个基于滞后值实现时序预测的TensorFlow模型，将其转换为PyTorch模型后，模型能正常运行但性能显著下降。以下是两个模型的代码，求分析两者间可能导致问题的明显差异，并给出排查建议。

TensorFlow模型代码 #

early_stop = EarlyStopping(monitor='val_loss',
min_delta=1e-3,
patience=3,
verbose=2, mode='auto')
tbCallBack = PlotLossesKeras()

model = Sequential()
model.add(LSTM(50, input_shape=(look_back, 1)))
model.add(Dropout(0.2))
model.add(Dense(1))
model.compile(loss='mean_squared_error', optimizer='adam')
model.fit(train_x, train_y,
          epochs=2000,
          batch_size=20, verbose=1)

PyTorch模型代码 #

class LSTMForecaster(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super().__init__()
        
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)
        self.dropout = nn.Dropout(0.2)
        self.linear = nn.Linear(hidden_size, output_size)
        
    def forward(self, input_seq):
        lstm_out, _ = self.lstm(input_seq)
        dropout = self.dropout(lstm_out)
        predictions = self.linear(dropout)
        return predictions
        
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model_1.parameters())
model_1 = LSTMForecaster(input_size=3, hidden_size=50, num_layers=1, output_size=1)

torch.manual_seed(45)
epochs = 2000 
batch_size = 20
num_batches = len(train_x) // batch_size
train_x = train_x.to(device)
test_x = test_x.to(device)
train_y = train_y.to(device)
test_y = test_y.to(device)

for epoch in range(epochs):
    for i in range(num_batches):
        # Get the current batch of data
        start = i * batch_size
        end = start + batch_size
        x_batch = train_x[start:end]
        y_batch = train_y[start:end]

        pred_y = model_1(x_batch)
        loss = loss_fn(pred_y, y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
            

    model_1.eval()

核心差异与排查建议 #

• 输入维度不匹配：TensorFlow中LSTM的input_shape=(look_back, 1)说明输入是(样本数, 时间步长, 特征数)，而PyTorch模型里input_size=3，和TensorFlow的特征数（1）完全不一致。需确保两者输入的特征数、时间步长维度对齐，同时注意PyTorch LSTM默认输入格式是(时间步长, 样本数, 特征数)，TensorFlow是(样本数, 时间步长, 特征数)，需检查数据维度是否转置正确。
• LSTM输出处理逻辑错误：TensorFlow的LSTM默认返回最后一个时间步的输出，而PyTorch的LSTM返回所有时间步的输出。你的PyTorch模型把所有时间步的输出传给全连接层，和TensorFlow逻辑完全不同。需修改forward函数，只取最后一个时间步的输出：lstm_out = lstm_out[-1, :, :]后再做后续处理。
• 优化器初始化顺序错误：PyTorch代码先初始化优化器再创建模型，导致优化器绑定的是未定义的模型参数，训练时根本没更新权重。必须先创建模型，再初始化优化器。
• 缺少早停机制：TensorFlow用EarlyStopping防止过拟合，PyTorch代码完全没实现，训练2000轮可能导致过拟合或训练不足。需手动实现早停逻辑，监控验证集损失，当损失不再下降时停止训练。
• 训练循环缺失验证逻辑：TensorFlow的fit默认处理验证集，PyTorch代码没有验证步骤，无法判断模型是否收敛。需在每个epoch结束后切换到eval模式计算验证损失，再切回train模式。
• Dropout模式切换错误：PyTorch的Dropout需要手动用model.train()和model.eval()切换模式，你的代码在epoch结束后调用model_1.eval()，但下一个epoch开始没切回model_1.train()，导致训练时Dropout不生效，影响拟合能力。
• 数据预处理/划分不一致：检查两个模型的训练集、验证集划分是否完全相同，数据归一化/标准化的方法、参数（如均值、标准差）是否一致，避免预处理差异导致性能下降。
• 批次处理逻辑问题：PyTorch代码丢弃了最后一个不足batch_size的样本，而TensorFlow的fit会处理剩余样本。需修改批次逻辑，确保所有训练样本都被用到。
• 权重初始化差异：TensorFlow和PyTorch的LSTM、全连接层默认初始化策略不同，可尝试手动对齐初始化方式，比如让PyTorch的权重初始化和TensorFlow一致。

内容的提问来源于stack exchange，提问作者Nickpick