Torch 2.0 Transformer 模块文本分类实战指南

以下是一个使用 Transformer 模块进行文本分类的示例：

首先，我们需要导入必要的库和模块：

import torch
import torch.nn as nn
import torch.optim as optim
from torchtext.datasets import AG_NEWS
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data import DataLoader
from torch.utils.data.dataset import random_split
from torch.nn.utils.rnn import pad_sequence
from torch.nn import TransformerEncoder, TransformerEncoderLayer

然后，我们需要定义一些超参数：

BATCH_SIZE = 64
EPOCHS = 5
MAX_LENGTH = 256
EMBED_DIM = 32
NHEAD = 4
DROPOUT = 0.2

接下来，我们需要加载数据集并进行预处理：

# 加载数据集
train_iter = AG_NEWS(split='train')
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>', '<bos>', '<eos>'])
vocab.set_default_index(vocab['<unk>'])

# 将文本转换为数值序列
def text_transform(text):
    tokens = torch.tensor([vocab[token] for token in tokenizer(text)], dtype=torch.long)
    return tokens

# 将标签转换为数值
label_transform = lambda x: int(x) - 1

# 创建数据集和数据加载器
def collate_batch(batch):
    label_list, text_list = [], []
    for (_label, _text) in batch:
        label_list.append(label_transform(_label))
        processed_text = text_transform(_text)
        text_list.append(processed_text)
    label_tensor = torch.tensor(label_list, dtype=torch.long)
    text_tensor = pad_sequence(text_list, padding_value=vocab['<pad>'], batch_first=True)
    return label_tensor, text_tensor

train_iter = AG_NEWS(split='train')
train_dataset = list(train_iter)
train_data, valid_data = random_split(train_dataset, [int(len(train_dataset)*0.95), len(train_dataset)-int(len(train_dataset)*0.95)])
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
valid_loader = DataLoader(valid_data, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)

现在，我们可以定义模型：

class TransformerModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, nhead, dropout):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        encoder_layers = TransformerEncoderLayer(embed_dim, nhead, dim_feedforward=256, dropout=dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layers, num_layers=2)
        self.fc = nn.Linear(embed_dim, 4)

    def forward(self, src):
        src = self.embedding(src)
        src = src.permute(1, 0, 2)
        output = self.transformer_encoder(src)
        output = output.mean(dim=0)
        output = self.fc(output)
        return output

在训练之前，我们需要定义损失函数和优化器：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = TransformerModel(len(vocab), EMBED_DIM, NHEAD, DROPOUT).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

最后，我们可以开始训练：

def train(model, iterator, optimizer, criterion):
    model.train()
    epoch_loss = 0
    for _, (labels, text) in enumerate(iterator):
        labels = labels.to(device)
        text = text.to(device)
        optimizer.zero_grad()
        output = model(text)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    return epoch_loss / len(iterator)

def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0
    with torch.no_grad():
        for _, (labels, text) in enumerate(iterator):
            labels = labels.to(device)
            text = text.to(device)
            output = model(text)
            loss = criterion(output, labels)
            epoch_loss += loss.item()
    return epoch_loss / len(iterator)

for epoch in range(EPOCHS):
    train_loss = train(model, train_loader, optimizer, criterion)
    valid_loss = evaluate(model, valid_loader, criterion)
    print('Epoch: %d | Train Loss: %.4f | Valid Loss: %.4f' % (epoch, train_loss, valid_loss))

这个示例演示了如何使用 Transformer 模块进行文本分类。在这个示例中，我们使用了 AG News 数据集，这是一个常用的新闻文本分类数据集。我们首先对文本进行了预处理，然后使用 Transformer 模块构建了一个模型，并使用 Adam 优化器进行训练。在训练过程中，我们使用交叉熵损失函数进行优化。最终，我们的模型在验证集上获得了较好的性能。