双向LSTM命名实体识别模型矩阵乘法报错问题求助

问题代码 #

class BiLSTMNERTagger(nn.Module):
    def __init__(self, emb_dim, hid_dim, n_layers, token_vocab_size, tag_vocab_size):
        super().__init__()

        self.embedding = nn.Embedding(token_vocab_size, emb_dim, padding_idx=0)
        self.rnn = nn.LSTM(emb_dim, hid_dim, num_layers=n_layers, batch_first=True, bidirectional=True)

        self.fc = nn.Linear(2*hid_dim, tag_vocab_size)

        self.softmax = nn.Softmax()

        self.dropout = nn.Dropout(0.1)
        self.hidden = None

    def forward(self, words, words_len):
      #YOUR CODE HERE
      out = self.dropout(self.embedding(words))
      out = nn.utils.rnn.pack_padded_sequence(out, words_len.cpu().numpy(), enforce_sorted=False, batch_first=True)
      out, (hidden, cell) = self.rnn(out)
      out = hidden[-1, :, :]
      self.hidden = hidden
      out = self.fc(out)
      prediction = self.softmax(out)
      return prediction

torch.manual_seed(42)
model = BiLSTMNERTagger(
    emb_dim=20,
    hid_dim=64,
    n_layers=2,
    token_vocab_size=len(tokens_vocab),
    tag_vocab_size=len(ner_vocab),)
model = model.to(device)

错误信息 #

115     def forward(self, input: Tensor) -> Tensor:
--> 116         return F.linear(input, self.weight, self.bias)
    117 
    118     def extra_repr(self) -> str:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x64 and 128x10)

错误原因 #

1. 双向LSTM隐藏状态结构误解：双向LSTM的hidden张量形状为(num_layers*2, batch_size, hid_dim)，前num_layers组是前向LSTM的隐藏层，后num_layers组是后向LSTM的隐藏层。你取hidden[-1, :, :]仅获取了最后一层的单方向隐藏状态，维度为(batch_size, hid_dim)，但全连接层fc期望输入维度是2*hid_dim，导致维度不匹配。
2. 任务逻辑偏差：NER是序列标注任务，需要为每个输入token输出对应实体标签，而非仅使用最后一个时间步的隐藏状态，当前逻辑错误地用了最终隐藏状态而非全序列输出。

修复方案 #

修改forward方法，正确处理双向LSTM的序列输出并恢复padding：

def forward(self, words, words_len):
    out = self.dropout(self.embedding(words))
    # 打包带padding的序列
    packed_out = nn.utils.rnn.pack_padded_sequence(out, words_len.cpu().numpy(), enforce_sorted=False, batch_first=True)
    # 经过双向LSTM
    packed_out, (hidden, cell) = self.rnn(packed_out)
    # 恢复padding，得到形状为(batch_size, seq_len, 2*hid_dim)的全序列输出
    out, _ = nn.utils.rnn.pad_packed_sequence(packed_out, batch_first=True)
    # 对每个时间步的输出做全连接变换
    out = self.fc(out)
    # 显式指定对标签维度做softmax，符合序列标注需求
    prediction = self.softmax(out, dim=-1)
    return prediction

额外优化提示 #

• 若后续使用CrossEntropyLoss，该损失函数已内置log softmax操作，可移除self.softmax以提升数值计算稳定性。

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