文本分類是自然語言處理（NLP）中的一個重要任務(wù)，廣泛應(yīng)用于情感分析、新聞分類、垃圾郵件檢測等領(lǐng)域。本教程將教你如何使用 TorchText 進(jìn)行文本分類。

一、導(dǎo)入必要的庫和模塊

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

二、加載數(shù)據(jù)集

TorchText 提供了多個預(yù)處理的文本分類數(shù)據(jù)集，如 AG_NEWS、SogouNews、DBpedia 等。我們將使用 AG_NEWS 數(shù)據(jù)集進(jìn)行演示。

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)

三、定義模型

我們將構(gòu)建一個簡單的文本分類模型，使用 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)

四、初始化模型和相關(guān)參數(shù)

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)

五、定義數(shù)據(jù)加載函數(shù)

為了處理不同長度的文本條目，我們使用自定義的 generate_batch 函數(shù)。

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

六、定義訓(xùn)練和評估函數(shù)

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_)

七、訓(xùn)練模型

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)')

九、使用模型進(jìn)行預(yù)測

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 進(jìn)行文本分類。在編程獅（W3Cschool）網(wǎng)站上，你可以找到更多關(guān)于 PyTorch 的詳細(xì)教程和實戰(zhàn)案例，幫助你進(jìn)一步提升深度學(xué)習(xí)技能，成為人工智能領(lǐng)域的編程大神。