PyTorch实战：三种神经网络模型对决IMDB电影评论情感分类，LSTM以88.76%准确率夺冠

情感分析是自然语言处理中的经典任务，IMDB电影评论数据集包含25000条训练数据和25000条测试数据，每条评论被标记为正面（1）或负面（0）。本文将使用PyTorch实现三种不同的神经网络模型来完成这个二分类任务。

数据集处理

1. 数据加载与预处理

我们使用torchtext库提供的IMDB数据集：

from torchtext import datasets
from torchtext.data import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator

train_pipe, test_pipe = datasets.IMDB()
tokenizer = get_tokenizer('basic_english')

2. 构建词表

vocab = build_vocab_from_iterator(
    [tokenizer(text) for label, text in train_pipe], 
    min_freq=2,
    specials=["<unk>", "<pad>"], 
    special_first=True
)
vocab.set_default_index(vocab["<unk>"])
padding_idx = vocab['<pad>']

3. 创建Dataset类

class IMDBDataset(Dataset):
    def __init__(self, pipe):
        self.texts = []
        self.labels = []
        for label, text in pipe:
            self.texts.append(text)
            self.labels.append(label)

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, index):
        text = self.texts[index]
        label = self.labels[index] - 1  # 将标签转换为0/1
        tokens = [vocab[token] for token in tokenizer(text)]
        return torch.tensor(tokens, dtype=torch.long), torch.tensor(label, dtype=torch.long)

4. 批处理与填充

max_sequence = 400

def collate_fn(batch):
    texts, labels = zip(*batch)
    texts = pad_sequence(texts, batch_first=True, padding_value=padding_idx)
    return texts[:, :max_sequence], torch.tensor(labels)

train_loader = DataLoader(train_dataset, batch_size=50, shuffle=True, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=50, shuffle=True, collate_fn=collate_fn)

模型实现

1. LSTM模型

双向LSTM能够捕捉上下文信息，适合处理序列数据：

class RNNModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(
            num_embeddings=vocab_len, 
            embedding_dim=100, 
            padding_idx=padding_idx
        )
        self.rnn = nn.LSTM(
            input_size=100, 
            hidden_size=256, 
            num_layers=2, 
            batch_first=True,
            dropout=0.5, 
            bidirectional=True
        )
        self.dropout = nn.Dropout()
        self.fc = nn.Linear(256 * 2, 2)

    def forward(self, x):
        out = self.embedding(x)
        out = self.dropout(out)
        _, (hidden, _) = self.rnn(out)
        out = torch.concat((hidden[-1], hidden[-2]), dim=1)
        out = self.dropout(out)
        return self.fc(out)

2. CNN模型

卷积神经网络通过多个卷积核提取局部特征：

class CNNModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(vocab_len, 100, padding_idx=padding_idx)
        
        self.cnn1 = nn.Sequential(
            nn.Conv1d(in_channels=100, out_channels=128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2)
        )
        self.cnn2 = nn.Sequential(
            nn.Conv1d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2, stride=2)
        )
        self.flatten = nn.Flatten()
        self.dropout = nn.Dropout(0.5)
        self.fc = nn.Linear(256 * 100, 2)

    def forward(self, x):
        out = self.embedding(x)
        out = out.permute(0, 2, 1)  # 转换维度以适应Conv1d
        out = self.cnn1(out)
        out = self.cnn2(out)
        out = self.dropout(out)
        out = self.flatten(out)
        return self.fc(out)

3. Transformer模型

使用自注意力机制捕获长距离依赖：

class TransformerModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(
            num_embeddings=vocab_len,
            embedding_dim=96,
            padding_idx=padding_idx,
        )
        self.positional_encoding = nn.Parameter(torch.zeros(1, max_sequence, 96))
        self.transformer = nn.Transformer(
            d_model=96,
            nhead=8,
            num_encoder_layers=2,
            num_decoder_layers=2,
            dim_feedforward=256,
            dropout=0.5,
            batch_first=True,
        )
        self.fc = nn.Linear(96, 2)

    def forward(self, x):
        seq_length = x.size(1)
        x = self.embedding(x) + self.positional_encoding[:, :seq_length, :]
        output = self.transformer(x, x)
        output = output.mean(dim=1)
        return self.fc(output)

训练流程

统一的训练循环：

model = RNNModel().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    for i, (texts, labels) in enumerate(train_loader):
        texts = texts.to(device)
        labels = labels.to(device)

        output = model(texts)
        loss = criterion(output, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

使用预训练词向量

使用GloVe预训练词向量可以提升模型性能：

from torchtext.vocab import GloVe

vocab = GloVe(name="6B", dim=100)

# 在Dataset中使用
tokens = [vocab.stoi[token] if token in vocab.stoi else 0 for token in tokenizer(text)]

实验结果对比

结果分析

LSTM表现最佳：在10个epoch时达到88.76%的测试准确率，说明双向LSTM能有效捕捉文本的上下文信息。

CNN性能稳定：CNN在5个epoch时达到峰值86.68%，训练速度快且不易过拟合。

Transformer表现中等：由于数据集规模有限，Transformer的优势未能充分发挥。

GloVe的影响：对于LSTM，GloVe预训练词向量在10个epoch时达到与训练词表相近的效果（88.76% vs 88.93%）。

总结

本文使用PyTorch实现了IMDB电影评论情感分类任务，对比了LSTM、CNN和Transformer三种模型的效果。实验表明，LSTM在该任务上表现最佳，准确率可达88.76%。完整代码可在GPU环境下运行，约需10-20分钟完成训练。

如需进一步提升性能，可以尝试：

• 调整超参数（学习率、批次大小、隐藏层维度）
• 使用更深的网络结构
• 集成多个模型
• 数据增强技术