从零开始：用PyTorch实现一个简易的文本分类器

前言

文本分类是自然语言处理（NLP）中最基础也最常用的任务之一。本文将带你从零开始，使用PyTorch实现一个基于神经网络的文本分类器。我们将使用IMDb电影评论数据集，判断评论是正面还是负面。

环境准备

确保已安装Python 3.8+和PyTorch（建议1.12以上版本）。

pip install torch torchvision torchaudio

步骤一：数据加载与预处理

我们使用torchtext库来加载数据。

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

# 加载数据集
train_iter, test_iter = IMDB(split=('train', 'test'))

# 定义分词器
tokenizer = get_tokenizer('basic_english')

# 构建词汇表
def yield_tokens(data_iter):
    for _, text in data_iter:
        yield tokenizer(text)

vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=['<unk>', '<pad>'])
vocab.set_default_index(vocab['<unk>'])

步骤二：创建数据加载器

将文本转换为数字序列，并生成批次。

from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence

def collate_batch(batch):
    label_list, text_list = [], []
    for (_label, _text) in batch:
        label_list.append(1 if _label == 'pos' else 0)
        text_tensor = torch.tensor(vocab(tokenizer(_text)), dtype=torch.long)
        text_list.append(text_tensor)
    labels = torch.tensor(label_list, dtype=torch.float32)
    texts = pad_sequence(text_list, batch_first=True, padding_value=vocab['<pad>'])
    return texts, labels

batch_size = 64
train_loader = DataLoader(list(train_iter), batch_size=batch_size, shuffle=True, collate_fn=collate_batch)
test_loader = DataLoader(list(test_iter), batch_size=batch_size, shuffle=False, collate_fn=collate_batch)

步骤三：定义模型

我们使用一个简单的嵌入层+循环神经网络（LSTM）+全连接层。

import torch.nn as nn

class TextClassifier(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_classes):
        super(TextClassifier, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=vocab['<pad>'])
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(hidden_dim * 2, num_classes)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        embedded = self.embedding(x)
        lstm_out, (hidden, cell) = self.lstm(embedded)
        # 取双向LSTM的最后隐藏状态
        hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)
        out = self.fc(hidden)
        return self.sigmoid(out).squeeze()

# 超参数
vocab_size = len(vocab)
embed_dim = 100
hidden_dim = 128
num_classes = 1

model = TextClassifier(vocab_size, embed_dim, hidden_dim, num_classes)

步骤四：训练模型

定义损失函数和优化器，开始训练。

import torch.optim as optim

criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for texts, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(texts)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')

步骤五：评估模型

在测试集上计算准确率。

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for texts, labels in test_loader:
        outputs = model(texts)
        predicted = (outputs > 0.5).float()
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total:.2f}%')

总结

恭喜你！你已经成功实现了一个文本分类器。你可以尝试调整超参数（如嵌入维度、隐藏层大小）或更换模型架构（如使用Transformer）来提升性能。

提示：完整代码可在GitHub仓库获取。