以下是一个基于PyTorch的文本分类模型的示例代码,用于将给定的文本分为多个预定义类别:
python
import torch
import torch.nn as nn
import torch.nn.functional as F
class TextClassifier(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, num_layers, bidirectional, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=num_layers, bidirectional=bidirectional, dropout=dropout)
self.fc = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, text, text_lengths):
embedded = self.dropout(self.embedding(text))
packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text_lengths.to('cpu'), enforce_sorted=False)
packed_output, (hidden, cell) = self.rnn(packed_embedded)
output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)
hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1) if self.rnn.bidirectional else hidden[-1,:,:])
return self.fc(hidden.squeeze(0))
该模型将输入的文本作为整数序列传递给嵌入层,然后通过多层LSTM层进行处理,最终输出每个类别的预测概率。
在训练模型之前,需要将文本序列转换为整数标记,通常使用分词器/标记器完成此任务。另外还需要定义优化器和损失函数来训练模型。
以下是一个完整的训练脚本的示例:
python
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 torchtext.data.utils import ngrams_iterator
from torchtext.data.utils import get_tokenizer
from torch.utils.data.dataset import random_split
from collections import Counter
# 获取数据集和分词器
train_iter = AG_NEWS(split='train')
tokenizer = get_tokenizer('basic_english')
# 构建词汇表
counter = Counter()
for (label, line) in train_iter:
counter.update(tokenizer(line))
vocab = build_vocab_from_iterator([counter])
vocab.set_default_index(vocab['<unk>'])
# 定义标记化函数和文本处理函数
def yield_tokens(data_iter):
for _, text in data_iter:
yield tokenizer(text)
def text_transform(tokenizer, vocab, data):
"""将文本数据转换为张量数据"""
data = [vocab[token] for token in tokenizer(data)]
return torch.tensor(data)
# 定义批次生成器
def collate_batch(batch):
label_list, text_list, offsets = [], [], [0]
for (_label, _text) in batch:
label_list.append(_label-1)
processed_text = torch.cat([text_transform(tokenizer, vocab, _text), torch.tensor([vocab['<eos>']])])
text_list.append(processed_text)
offsets.append(processed_text.size(0))
label_list = torch.tensor(label_list)
offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
text_list = torch.cat(text_list)
return label_list, text_list, offsets
# 构建数据集和数据加载器
train_iter, test_iter = AG_NEWS()
train_iter = list(train_iter)
test_iter = list(test_iter)
train_dataset = list(map(lambda x: (x[0], x[1]), train_iter))
test_dataset = list(map(lambda x: (x[0], x[1]), test_iter))
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, collate_fn=collate_batch)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True, collate_fn=collate_batch)
# 创建模型和优化器
model = TextClassifier(len(vocab), 64, 128, 4, 2, True, 0.5)
optimizer = optim.Adam(model.parameters())
# 定义损失函数和训练函数
criterion = nn.CrossEntropyLoss()
def train(model, iterator, optimizer, criterion):
epoch_loss = 0
model.train()
for (label, text, offsets) in iterator:
optimizer.zero_grad()
predictions = model(text, offsets)
loss = criterion(predictions, label)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
return epoch_loss / len(iterator)
# 训练模型
N_EPOCHS = 10
for epoch in range(N_EPOCHS):
train_loss = train(model, train_loader, optimizer, criterion)
print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')
在训练过程结束后,可以使用该模型对新的文本进行分类。具体方法是将文本转换为整数标记序列,然后使用模型进行预测:
python
# 对新文本进行分类
def predict(model, sentence):
model.eval()
tokenized = torch.tensor([vocab[token] for token in tokenizer(sentence)])
length = torch.tensor([len(tokenized)])
prediction = model(tokenized, length)
return F.softmax(prediction, dim=1).detach().numpy()[0]
# 进行预测
test_sentence = "World markets are reacting to the news that the UK is set to leave the European Union."
pred_probs = predict(model, test_sentence)
print(pred_probs)
以上代码示例中使用了AG_NEWS数据集作为示例训练数据,可通过以下方式加载数据集:
python
from torchtext.datasets import AG_NEWS
train_iter = AG_NEWS(split='train')
test_iter = AG_NEWS(split='test')
该数据集包含四个类别的新闻数据,每个类别各有120,000个训练示例和7,600个测试示例。完整的训练脚本和数据集可以在PyTorch官方文档中找到。