将PyTorch单向LSTM改为双向时遇张量维度不匹配错误求解决

问题根源 #

报错RuntimeError: The size of tensor a (1534) must match the size of tensor b (767)的核心原因是双向LSTM的输出维度处理逻辑错误：

• 双向LSTM的输出特征数是hidden_size * 2，但原模型的线性层仅按单方向hidden_size设置输入维度，导致后续维度混乱
• 训练时手动取了单方向输出output[:,0,:]匹配标签，但测试时直接展开所有输出，导致输出长度是标签的2倍，触发维度不匹配

修复方案 #

1. 修改模型类代码

class RNN(nn.Module):
    def __init__(self, i_size, h_size, n_layers, o_size, dropout=0.1, bidirectional=True):
        super(RNN, self).__init__()
        # 显式计算方向数，替代布尔值算术运算，提升可读性
        self.num_directions = 2 if bidirectional else 1
        self.rnn = nn.LSTM(
            input_size=i_size,
            hidden_size=h_size,
            num_layers=n_layers,
            dropout=dropout,
            bidirectional=bidirectional
        )
        # 双向LSTM输出特征数为hidden_size * 方向数，线性层需适配该维度
        self.out = nn.Linear(h_size * self.num_directions, o_size)

    def forward(self, x, h_state):
        r_out, hidden_state = self.rnn(x, h_state)
        # 直接使用双向拼接后的特征，无需手动拆分方向维度
        outs = self.out(r_out)
        return outs, hidden_state

2. 统一训练与测试的维度处理逻辑

history = []
x_test_cuda = torch.tensor(x_test).float().cuda()
y_test_cuda = torch.tensor(y_test).float().cuda()

inputs_cuda = torch.tensor(x_train).float().cuda()
labels_cuda = torch.tensor(y_train).float().cuda()

for epoch in range(num_epochs):
    rnn.train()
    output, _ = rnn(inputs_cuda, hidden_state)
    # 统一将输出展开为一维，与标签维度匹配
    loss = criterion(output.view(-1), labels_cuda)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 20 == 0:
        rnn.eval()
        test_output, _ = rnn(x_test_cuda, hidden_state)
        # 测试时使用与训练完全一致的维度处理方式
        test_loss = criterion(test_output.view(-1), y_test_cuda)
        print(f'epoch {epoch}, loss {loss.item()}, eval loss {test_loss.item()}')
    else:
        print(f'epoch {epoch}, loss {loss.item()}')
    
    history.append(loss.item())

关键修改说明 #

• 线性层维度适配：将线性层输入维度改为hidden_size * num_directions，直接利用双向LSTM拼接后的特征进行预测
• 简化forward逻辑：移除手动拆分方向的代码，避免维度处理失误
• 统一维度匹配：训练和测试阶段均用view(-1)展开输出，确保与标签的一维结构完全对齐

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