PyTorch中LSTM分类模型实现详解

2026年1月7日 互联网

PyTorch中LSTM分类模型实现详解

LSTM（长短期记忆网络）作为循环神经网络（RNN）的改进版本，在处理序列数据时展现出强大的特征提取能力，尤其适用于文本分类、时间序列预测等任务。本文将通过PyTorch框架实现一个完整的LSTM分类模型，从数据预处理到模型部署进行系统性讲解。

一、LSTM分类模型核心原理

LSTM通过引入门控机制（输入门、遗忘门、输出门）解决了传统RNN的梯度消失问题，能够捕捉长距离依赖关系。在分类任务中，LSTM层将输入序列编码为固定维度的上下文向量，再通过全连接层输出分类结果。

典型架构包含：

  • 嵌入层（Embedding）：将离散词索引映射为密集向量
  • LSTM层：提取序列特征，输出每个时间步的隐藏状态
  • 池化层：通常取最后一个时间步的隐藏状态或所有时间步的平均值
  • 全连接层：完成类别预测

二、数据预处理关键步骤

1. 文本序列化处理

  1. from torchtext.data import Field, TabularDataset
  2. from torchtext.vocab import Vectors
  4. # 定义文本处理字段
  5. TEXT = Field(tokenize='spacy', lower=True, include_lengths=True)
  6. LABEL = Field(sequential=False, use_vocab=False)
  8. # 加载数据集（示例为CSV格式）
  9. train_data, test_data = TabularDataset.splits(
  10.     path='./data',
  11.     train='train.csv',
  12.     test='test.csv',
  13.     format='csv',
  14.     fields=[('text', TEXT), ('label', LABEL)],
  15.     skip_header=True
  16. )

2. 构建词汇表与数值化

  1. MAX_VOCAB_SIZE = 25000
  2. TEXT.build_vocab(train_data, max_size=MAX_VOCAB_SIZE, vectors="glove.6B.100d")
  3. LABEL.build_vocab(train_data)

3. 创建迭代器

  1. BATCH_SIZE = 64
  2. train_iterator, test_iterator = BucketIterator.splits(
  3.     (train_data, test_data),
  4.     batch_size=BATCH_SIZE,
  5.     sort_within_batch=True,
  6.     sort_key=lambda x: len(x.text),
  7.     device=device
  8. )

三、LSTM模型实现代码

1. 基础LSTM分类模型

  1. import torch.nn as nn
  3. class LSTMClassifier(nn.Module):
  4.     def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, dropout):
  5.         super().__init__()
  6.         self.embedding = nn.Embedding(vocab_size, embedding_dim)
  7.         self.lstm = nn.LSTM(embedding_dim, 
  8.                            hidden_dim, 
  9.                            num_layers=n_layers,
  10.                            dropout=dropout if n_layers > 1 else 0)
  11.         self.fc = nn.Linear(hidden_dim, output_dim)
  12.         self.dropout = nn.Dropout(dropout)
  14.     def forward(self, text, text_lengths):
  15.         # text: [sent len, batch size]
  16.         embedded = self.dropout(self.embedding(text))
  17.         # embedded: [sent len, batch size, emb dim]
  19.         packed_embedded = nn.utils.rnn.pack_padded_sequence(
  20.             embedded, text_lengths.to('cpu'))
  22.         packed_output, (hidden, cell) = self.lstm(packed_embedded)
  23.         # hidden: [num layers, batch size, hid dim]
  25.         hidden = self.dropout(hidden[-1,:,:])
  26.         # hidden: [batch size, hid dim]
  28.         return self.fc(hidden)

2. 双向LSTM改进实现

  1. class BiLSTMClassifier(nn.Module):
  2.     def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, dropout):
  3.         super().__init__()
  4.         self.embedding = nn.Embedding(vocab_size, embedding_dim)
  5.         self.lstm = nn.LSTM(embedding_dim,
  6.                            hidden_dim // 2,  # 双向时隐藏维度减半
  7.                            num_layers=n_layers,
  8.                            bidirectional=True,
  9.                            dropout=dropout if n_layers > 1 else 0)
  10.         self.fc = nn.Linear(hidden_dim, output_dim)
  11.         self.dropout = nn.Dropout(dropout)
  13.     def forward(self, text, text_lengths):
  14.         embedded = self.dropout(self.embedding(text))
  15.         packed_embedded = nn.utils.rnn.pack_padded_sequence(
  16.             embedded, text_lengths.to('cpu'))
  18.         packed_output, (hidden, cell) = self.lstm(packed_embedded)
  19.         # 双向LSTM的hidden拼接
  20.         hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1))
  22.         return self.fc(hidden)

四、模型训练与评估

1. 训练循环实现

  1. def train(model, iterator, optimizer, criterion, device):
  2.     epoch_loss = 0
  3.     epoch_acc = 0
  4.     model.train()
  6.     for batch in iterator:
  7.         optimizer.zero_grad()
  8.         text, text_lengths = batch.text
  9.         predictions = model(text, text_lengths).squeeze(1)
  10.         loss = criterion(predictions, batch.label)
  11.         acc = categorical_accuracy(predictions, batch.label)
  13.         loss.backward()
  14.         optimizer.step()
  16.         epoch_loss += loss.item()
  17.         epoch_acc += acc.item()
  19.     return epoch_loss / len(iterator), epoch_acc / len(iterator)

2. 评估函数实现

  1. def evaluate(model, iterator, criterion, device):
  2.     epoch_loss = 0
  3.     epoch_acc = 0
  4.     model.eval()
  6.     with torch.no_grad():
  7.         for batch in iterator:
  8.             text, text_lengths = batch.text
  9.             predictions = model(text, text_lengths).squeeze(1)
  10.             loss = criterion(predictions, batch.label)
  11.             acc = categorical_accuracy(predictions, batch.label)
  13.             epoch_loss += loss.item()
  14.             epoch_acc += acc.item()
  16.     return epoch_loss / len(iterator), epoch_acc / len(iterator)

五、关键参数优化建议

  1. 超参数调优：

    • 隐藏层维度：通常设置128-512，根据任务复杂度调整
    • LSTM层数：1-3层，深层网络需要配合残差连接
    • Dropout率：0.2-0.5之间，防止过拟合

  2. 性能优化技巧：

    • 使用预训练词向量（如GloVe、Word2Vec）
    • 采用梯度裁剪（gradient clipping）防止梯度爆炸
    • 学习率调度：使用ReduceLROnPlateau动态调整

  3. 常见问题处理：

    • 梯度消失：改用GRU或增加LSTM单元数
    • 过拟合：增加Dropout层，使用L2正则化
    • 内存不足：减小batch size，使用梯度累积

六、完整训练流程示例

  1. import torch
  2. from torch.optim import Adam
  4. # 初始化模型
  5. INPUT_DIM = len(TEXT.vocab)
  6. EMBEDDING_DIM = 100
  7. HIDDEN_DIM = 256
  8. OUTPUT_DIM = len(LABEL.vocab)
  9. N_LAYERS = 2
  10. DROPOUT = 0.5
  12. model = LSTMClassifier(INPUT_DIM, EMBEDDING_DIM, HIDDEN_DIM, OUTPUT_DIM, N_LAYERS, DROPOUT)
  13. optimizer = Adam(model.parameters())
  14. criterion = nn.CrossEntropyLoss()
  15. model = model.to(device)
  16. criterion = criterion.to(device)
  18. # 训练循环
  19. N_EPOCHS = 10
  20. for epoch in range(N_EPOCHS):
  21.     train_loss, train_acc = train(model, train_iterator, optimizer, criterion, device)
  22.     valid_loss, valid_acc = evaluate(model, valid_iterator, criterion, device)
  24.     print(f'Epoch: {epoch+1:02}')
  25.     print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')
  26.     print(f'\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')

七、部署与扩展建议

  1. 模型导出：

    1. torch.save(model.state_dict(), 'lstm_classifier.pt')
    2. # 加载模型示例
    3. loaded_model = LSTMClassifier(INPUT_DIM, EMBEDDING_DIM, HIDDEN_DIM, OUTPUT_DIM, N_LAYERS, DROPOUT)
    4. loaded_model.load_state_dict(torch.load('lstm_classifier.pt'))

  2. 服务化部署：

    • 使用TorchScript转换模型
    • 部署为REST API服务
    • 容器化部署（Docker + Kubernetes）

  3. 进阶方向：

    • 结合注意力机制（Attention）
    • 尝试Transformer架构对比
    • 多模态融合（结合图像特征）

通过系统化的实现和优化，LSTM分类模型在文本分类任务中可以达到90%以上的准确率。实际开发中，建议从简单架构开始，逐步增加复杂度，同时密切关注验证集指标变化，避免过度拟合。对于大规模数据集，可考虑使用分布式训练框架加速模型收敛。

最新文章