PyTorch 使用 TorchText 进行文本分类

文本分类是自然语言处理（NLP）中的一个重要任务，广泛应用于情感分析、新闻分类、垃圾邮件检测等领域。本教程将教你如何使用 TorchText 进行文本分类。

一、导入必要的库和模块

import torch
import torchtext
from torchtext.datasets import text_classification
import os

二、加载数据集

TorchText 提供了多个预处理的文本分类数据集，如 AG_NEWS、SogouNews、DBpedia 等。我们将使用 AG_NEWS 数据集进行演示。

NGRAMS = 2  # 使用二元语法
if not os.path.isdir('./.data'):
    os.mkdir('./.data')
train_dataset, test_dataset = text_classification.DATASETS['AG_NEWS'](
    root='./.data', ngrams=NGRAMS, vocab=None)

三、定义模型

我们将构建一个简单的文本分类模型，使用 EmbeddingBag 层和线性层。

import torch.nn as nn
import torch.nn.functional as F


class TextSentiment(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super().__init__()
        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, sparse=True)
        self.fc = nn.Linear(embed_dim, num_class)
        self.init_weights()


    def init_weights(self):
        initrange = 0.5
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()


    def forward(self, text, offsets):
        embedded = self.embedding(text, offsets)
        return self.fc(embedded)

四、初始化模型和相关参数

VOCAB_SIZE = len(train_dataset.get_vocab())
EMBED_DIM = 32
NUM_CLASS = len(train_dataset.get_labels())
BATCH_SIZE = 16


model = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUM_CLASS).to(device)

五、定义数据加载函数

为了处理不同长度的文本条目，我们使用自定义的 generate_batch 函数。

def generate_batch(batch):
    label = torch.tensor([entry[0] for entry in batch])
    text = [entry[1] for entry in batch]
    offsets = [0] + [len(entry) for entry in text]
    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
    text = torch.cat(text)
    return text, offsets, label

六、定义训练和评估函数

from torch.utils.data import DataLoader


def train_func(sub_train_):
    train_loss = 0
    train_acc = 0
    data = DataLoader(sub_train_, batch_size=BATCH_SIZE, shuffle=True, collate_fn=generate_batch)
    for i, (text, offsets, cls) in enumerate(data):
        optimizer.zero_grad()
        text, offsets, cls = text.to(device), offsets.to(device), cls.to(device)
        output = model(text, offsets)
        loss = criterion(output, cls)
        train_loss += loss.item()
        loss.backward()
        optimizer.step()
        train_acc += (output.argmax(1) == cls).sum().item()
    scheduler.step()
    return train_loss / len(sub_train_), train_acc / len(sub_train_)


def test(data_):
    loss = 0
    acc = 0
    data = DataLoader(data_, batch_size=BATCH_SIZE, collate_fn=generate_batch)
    for text, offsets, cls in data:
        text, offsets, cls = text.to(device), offsets.to(device), cls.to(device)
        with torch.no_grad():
            output = model(text, offsets)
            loss = criterion(output, cls)
            loss += loss.item()
            acc += (output.argmax(1) == cls).sum().item()
    return loss / len(data_), acc / len(data_)

七、训练模型

import time
from torch.utils.data.dataset import random_split


N_EPOCHS = 5
min_valid_loss = float('inf')


criterion = torch.nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=4.0)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.9)


train_len = int(len(train_dataset) * 0.95)
sub_train_, sub_valid_ = random_split(train_dataset, [train_len, len(train_dataset) - train_len])


for epoch in range(N_EPOCHS):
    start_time = time.time()
    train_loss, train_acc = train_func(sub_train_)
    valid_loss, valid_acc = test(sub_valid_)
    secs = int(time.time() - start_time)
    mins = secs // 60
    secs = secs % 60
    print(f'Epoch: {epoch + 1}, Time: {mins}m {secs}s')
    print(f'\tLoss: {train_loss:.4f}(train)\t|\tAcc: {train_acc * 100:.1f}%(train)')
    print(f'\tLoss: {valid_loss:.4f}(valid)\t|\tAcc: {valid_acc * 100:.1f}%(valid)')

八、评估模型

print('Checking the results of test dataset...')
test_loss, test_acc = test(test_dataset)
print(f'\tLoss: {test_loss:.4f}(test)\t|\tAcc: {test_acc * 100:.1f}%(test)')

九、使用模型进行预测

import re
from torchtext.data.utils import ngrams_iterator
from torchtext.data.utils import get_tokenizer


ag_news_label = {1: "World", 2: "Sports", 3: "Business", 4: "Sci/Tec"}


def predict(text, model, vocab, ngrams):
    tokenizer = get_tokenizer("basic_english")
    with torch.no_grad():
        text = torch.tensor([vocab[token] for token in ngrams_iterator(tokenizer(text), ngrams)])
        output = model(text, torch.tensor([0]))
        return output.argmax(1).item() + 1


vocab = train_dataset.get_vocab()
model = model.to("cpu")


ex_text_str = "MEMPHIS, Tenn. – Four days ago, Jon Rahm was enduring the season's worst weather conditions on Sunday at The Open on his way to a closing 75 at Royal Portrush, which considering the wind and the rain was a respectable showing. Thursday's first round at the WGC-FedEx St. Jude Invitational was another story. With temperatures in the mid-80s and hardly any wind, the Spaniard was 13 strokes better in a flawless round. Thanks to his best putting performance on the PGA Tour, Rahm finished with an 8-under 62 for a three-stroke lead, which was even more impressive considering he'd never played the front nine at TPC Southwind."


print("This is a %s news" % ag_news_label[predict(ex_text_str, model, vocab, 2)])

通过本教程，你掌握了如何使用 PyTorch 和 TorchText 进行文本分类。在编程狮（W3Cschool）网站上，你可以找到更多关于 PyTorch 的详细教程和实战案例，帮助你进一步提升深度学习技能，成为人工智能领域的编程大神。