PyTorch LSTM模型生成输出时hidden[0]尺寸不匹配报错求助

错误信息 #

Traceback (most recent call last):
  File "d:\Projects\Crptic\CrpticMain\TrainedModel.py", line 146, in <module>
    print(generate_code(model, input_text.lower()))
  File "d:\Projects\Crptic\CrpticMain\TrainedModel.py", line 115, in generate_code
    output, hidden = model(input_tensor, hidden)
  File "C:\Users\aadi_\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "d:\Projects\Crptic\CrpticMain\TrainedModel.py", line 38, in forward
    output, hidden = self.lstm(input, hidden)
  File "C:\Users\aadi_\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\aadi_\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\rnn.py", line 810, in forward
    self.check_forward_args(input, hx, batch_sizes)
  File "C:\Users\aadi_\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\rnn.py", line 731, in check_forward_args
    self.check_hidden_size(hidden[0], self.get_expected_hidden_size(input, batch_sizes),
  File "C:\Users\aadi_\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\rnn.py", line 239, in check_hidden_size
    raise RuntimeError(msg.format(expected_hidden_size, list(hx.size())))
RuntimeError: Expected hidden[0] size (1, 16, 256), got [1, 1, 256]

问题背景 #

使用PyTorch构建的LSTM模型在生成输出时持续抛出上述隐藏层尺寸不匹配错误，调整生成函数和模型对象后仍未解决。模型在CPU上训练，数据集规模较小，期望输入"make a print statement"时输出"print('hello world')"，相关词汇均在训练词表中。

核心错误原因 #

1. 输入维度不统一：训练时逐个传入单个字符（维度为标量），生成时传入整段序列，forward方法中随意使用view调整形状，导致LSTM对隐藏层的batch_size期望不一致。错误中期望的(1,16,256)里的16是生成时输入序列的长度，而初始化的隐藏层batch_size为1，两者不匹配。
2. 词表索引不一致：训练时用字符的ASCII码（ord(c)）作为张量索引，生成时却临时基于输入文本创建自定义词表，导致嵌入层输入的索引值完全不对应训练时的权重，属于逻辑错误。
3. Forward方法维度处理混乱：未遵循LSTM输入的标准格式（(seq_len, batch_size)），随意调整张量形状导致输入与隐藏层的维度不兼容。

修正方案及代码 #

1. 修正模型Forward方法 #

明确输入维度为(seq_len, batch_size)，遵循LSTM的标准输入格式：

class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.encoder = nn.Embedding(input_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size)
        self.decoder = nn.Linear(hidden_size, output_size)

    def forward(self, inp, hidden):
        # inp shape: (seq_len, batch_size)
        embedded = self.encoder(inp)  # shape: (seq_len, batch_size, hidden_size)
        output, hidden = self.lstm(embedded, hidden)  # output shape: (seq_len, batch_size, hidden_size)
        output = self.decoder(output)  # shape: (seq_len, batch_size, output_size)
        return output, hidden

    def init_hidden(self, batch_size):
        return (torch.zeros(1, batch_size, self.hidden_size),
                torch.zeros(1, batch_size, self.hidden_size))

2. 修正训练函数 #

统一输入维度，避免逐个字符传入，保证与生成时的输入格式一致：

def train(model, optimizer, criterion, input_tensor, target_tensor):
    batch_size = 1
    hidden = model.init_hidden(batch_size)
    optimizer.zero_grad()
    
    # 调整维度为(seq_len, batch_size)
    input_tensor = input_tensor.unsqueeze(1)  # shape: (input_len, 1)
    target_tensor = target_tensor.unsqueeze(1)  # shape: (target_len, 1)
    
    # 先处理输入序列
    output, hidden = model(input_tensor, hidden)
    
    loss = 0
    # 按原逻辑处理目标序列（teacher forcing）
    for i in range(target_tensor.size(0)):
        step_input = target_tensor[i:i+1]  # shape: (1, 1)
        output, hidden = model(step_input, hidden)
        loss += criterion(output[0], target_tensor[i])
    
    loss.backward()
    optimizer.step()
    return loss.item() / target_tensor.size(0)

3. 修正生成函数 #

统一词表索引规则（与训练时一致用ASCII码），并调整输入维度：

def generate_code(model, input_text):
    with torch.no_grad():
        # 沿用训练时的索引规则：字符ASCII码
        input_chars = list(input_text.lower())
        input_tensor = torch.tensor([ord(c) for c in input_chars], dtype=torch.long)
        # 调整维度为(seq_len, batch_size)
        input_tensor = input_tensor.unsqueeze(1)  # shape: (input_len, 1)
        
        # 初始化隐藏层，batch_size与输入一致
        hidden = model.init_hidden(1)
        # 处理输入序列
        output, hidden = model(input_tensor, hidden)
        
        # 从最后一步输出开始生成
        _, topi = output[-1].topk(1)
        current_char = chr(topi.item())
        code = current_char
        
        # 生成直到终止符
        while current_char not in ['\n', ':']:
            # 调整输入维度为(1, 1)
            input_tensor = torch.tensor([ord(current_char)], dtype=torch.long).unsqueeze(1)
            output, hidden = model(input_tensor, hidden)
            _, topi = output[-1].topk(1)
            current_char = chr(topi.item())
            code += current_char
        
        return code

4. 修正主函数 #

添加模型评估模式，避免训练时的Dropout等影响生成：

if __name__ == "__main__":
    model = LSTMModel(input_size, hidden_size, output_size)
    model.load_state_dict(torch.load(model_path))
    model.eval()  # 进入评估模式
    
    input_text = "make a print statement"
    print('here is what is being fed: ' + input_text.lower())
    print(generate_code(model, input_text.lower()))

额外注意事项 #

• 训练时的input_size设为256，对应ASCII码的范围（0-255），这个设置是合理的，无需修改。
• 如果训练时模型未正确保存，需要重新训练修正后的模型，再进行生成测试。

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