MPS环境下测试集Batch Size影响LSTM模型性能异常问题

问题背景 #

模型结构：

Model(
     (lstm): LSTM(3, 32, num_layers=3, batch_first=True, dropout=0.7)  
     (dense): Linear(in_features=32, out_features=2, bias=True) 
)

训练环境：Apple M2芯片的MPS环境，损失函数为带类别权重的Cross Entropy Loss，优化器AdamW，用sklearn.classification_report的F-Score评估。

异常现象：训练集和测试集batch size均为64时性能正常；测试集batch size改为256后，性能大幅下降且无提升。已正确设置train()/eval()模式。

可能的原因与解决方案 #

1. LSTM隐藏状态未正确重置（最高优先级） #

PyTorch的LSTM默认会在每次forward时初始化全0的隐藏状态，但MPS后端部分版本存在大batch下隐藏状态复用的bug，导致后续batch预测依赖前序batch的状态，累积错误引发性能暴跌。

解决方法：
手动在每个测试batch前初始化LSTM的隐藏状态：

# 定义隐藏状态初始化函数
def init_lstm_hidden(batch_size, hidden_size, num_layers, device):
    h0 = torch.zeros(num_layers, batch_size, hidden_size, device=device)
    c0 = torch.zeros(num_layers, batch_size, hidden_size, device=device)
    return (h0, c0)

# 修改测试代码的batch循环
with torch.no_grad():
    for batch_id, (X, y) in enumerate(test_dataloader):
        X, y = X.to(device), y.to(device)
        batch_size = X.shape[0]
        # 初始化当前batch的隐藏状态
        hidden = init_lstm_hidden(batch_size, 32, 3, device)
        # 手动调用LSTM和全连接层，传入隐藏状态
        lstm_out, _ = model.lstm(X, hidden)
        # 假设是序列分类任务，取最后一个时间步的输出
        y_pred = model.dense(lstm_out[:, -1, :])
        y_pred = torch.argmax(y_pred, dim=1)
        
        y_pred_all.append(y_pred.cpu().numpy())
        y_all.append(y.cpu().numpy())
        bar.update(batch_id)

注：若你的模型forward方法已封装LSTM和全连接层的逻辑，需修改forward以支持传入隐藏状态参数。

2. MPS后端LSTM Dropout禁用不彻底 #

model.eval()理论上会关闭Dropout，但MPS后端在大batch下可能未正确禁用LSTM的层间Dropout（LSTM的dropout参数作用于层与层之间）。

解决方法：
测试时手动强制关闭Dropout：

model.eval()
# 遍历模型模块，关闭LSTM的dropout
for module in model.modules():
    if isinstance(module, torch.nn.LSTM):
        module.dropout = 0.0

或升级PyTorch至最新稳定版本，修复MPS后端的Dropout实现bug。

3. MPS数值精度偏差 #

大batch下MPS的float32计算可能存在精度偏差，导致模型输出的logit值偏移，影响最终分类结果。

解决方法：

• 临时切换到CPU运行测试，验证性能是否恢复：

accuracy, f_score = test_model(model, test_dataloader, device="cpu")

若CPU上大batch性能正常，可尝试：

1. 强制模型使用float32精度：

model = model.float()

2. 升级PyTorch到支持MPS的最新版本，修复精度相关问题。

4. 数据加载逻辑不一致 #

检查测试集DataLoader的设置，确保大batch与小batch的处理逻辑完全一致：

• 确认测试集shuffle=False，避免数据分布变化
• 若存在基于batch的动态预处理（如batch归一化），改用全局统计量进行归一化，避免大batch下统计量偏移。

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