PyTorch RNN模型运行报错输入需3维却得到2维问题排查

PyTorch 内置的 RNN、LSTM、GRU 等循环神经网络层，要求输入张量必须为3维，默认维度顺序为(序列长度, batch大小, 特征维度)，若设置batch_first=True则维度顺序为(batch大小, 序列长度, 特征维度)，触发该报错说明输入循环层的张量仅存在2个维度，缺少序列长度或特征维度，可按以下步骤排查解决：

• 校验数据输入维度：不要仅验证数据集返回的单样本维度，在 Lightning 模型的training_step、validation_step方法最开头打印输入张量的shape，确认Dataloader打包后、送入模型前的维度是否符合3维要求。如果是单特征序列数据，常见问题是预处理时丢失了最后一维特征维度，可通过x = x.unsqueeze(-1)补全维度。
• 核对RNN层配置与输入维度的匹配性：如果你的输入维度是(batch大小, 序列长度, 特征维度)，必须给RNN层传入batch_first=True参数，否则PyTorch会默认将第一个维度识别为序列长度，出现维度匹配错误。
• 检查forward方法前置处理逻辑：确认在将张量送入RNN层前，没有执行view、flatten等操作误将3维张量压缩为2维，若存在全连接层和RNN层的衔接，需注意维度转换的顺序。
• 单步推理场景维度补全：若你是在逐时间步推理的场景下触发报错，单步输入的2维张量(batch大小, 特征维度)需要补全序列长度维度，batch_first模式下可通过x = x.unsqueeze(1)调整为3维。

正确实现参考代码 #

import torch
from torch import nn

# 模拟错误输入：batch大小32，序列长度10，单特征，仅2维
x = torch.randn(32, 10)
print(x.shape) # torch.Size([32, 10]) 2维输入会触发报错

# 补全特征维度，转为3维
x = x.unsqueeze(-1)
print(x.shape) # torch.Size([32, 10, 1]) 符合batch_first模式的3维要求

# 初始化RNN层，配置batch_first匹配输入维度
rnn_layer = nn.RNN(input_size=1, hidden_size=64, batch_first=True)
output, hidden = rnn_layer(x)

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