博客

Pytorch从0实现Transformer的实践

  • Post category:Python

PyTorch是目前非常流行的机器学习框架之一，它支持众多的模型实现，包括用于处理序列数据的Transformer模型。下面是从零开始实现Transformer模型的步骤。

1. 准备数据

一般来说，Transformer模型需要输入一个序列$x=(x_1, x_2, …, x_n)$，并输出另一个序列$y=(y_1, y_2, …, y_m)$。为了方便训练，我们需要将这些序列转换为数字，最简单的方法是将每个单词映射到一个唯一的整数，并通过padding使每个输入序列具有相同的长度。

示例代码：

import torchtext

# 定义数据集的字段
TEXT = torchtext.data.Field(tokenize='spacy', lower=True)
LABEL = torchtext.data.LabelField(dtype=torch.float)

# 构建数据集
train_data, test_data = torchtext.datasets.IMDB.splits(TEXT, LABEL)

# 构建词汇表
TEXT.build_vocab(train_data, max_size=10000)
LABEL.build_vocab(train_data)

在上面的示例中，我们使用了torchtext库来加载IMDB数据集，并定义了用于处理文本和标签的Field对象。然后，我们可以通过调用build_vocab方法来构建词汇表，并将每个单词映射到一个整数。

2. 构建模型

接下来，我们需要定义Transformer模型的结构。下面是一个简单的Transformer模型实现示例：

import torch
import torch.nn as nn

class TransformerModel(nn.Module):

    def __init__(self, ntoken, ninp, nhead, nhid, nlayers, dropout=0.5):
        super(TransformerModel, self).__init__()
        from torch.nn import TransformerEncoder, TransformerEncoderLayer
        self.model_type = 'Transformer'
        self.src_mask = None
        self.pos_encoder = PositionalEncoding(ninp, dropout)
        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
        self.encoder = nn.Embedding(ntoken, ninp)
        self.ninp = ninp
        self.decoder = nn.Linear(ninp, ntoken)

        self.init_weights()

    def _generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
        return mask

    def init_weights(self):
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, src):
        if self.src_mask is None or self.src_mask.size(0) != len(src):
            device = src.device
            mask = self._generate_square_subsequent_mask(len(src)).to(device)
            self.src_mask = mask

        src = self.encoder(src) * math.sqrt(self.ninp)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src, self.src_mask)
        output = self.decoder(output)
        return output

在这个示例中，我们使用了nn.TransformerEncoder和nn.MultiheadAttention来实现Transformer模型。我们还定义了一个PositionalEncoding类来将位置信息嵌入到词嵌入中。

3. 定义损失函数和优化器

定义损失函数和优化器是训练模型的重要一步。在这个示例中，我们使用交叉熵损失和Adam优化器：

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

其中nn.CrossEntropyLoss()是交叉熵损失函数的实现，而torch.optim.Adam()是Adam优化器的实现。

4. 训练模型

在数据准备、模型构建和优化器定义完成后，我们可以开始训练模型了。下面是一个简单的模型训练示例：

def train(model, iterator, optimizer, criterion):
    model.train()
    epoch_loss = 0
    for batch in iterator:
        optimizer.zero_grad()
        src = batch.text
        trg = batch.label
        output = model(src)
        loss = criterion(output.view(-1, output.shape[-1]), trg)
        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_iterator, optimizer, criterion)
    print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')

在这个示例中，我们循环遍历训练数据集，并在每个批次完成后更新模型参数。最后，我们计算每个epoch的平均损失并输出。

5. 应用模型

训练完模型后，我们可以将其应用到新的数据上。下面是一个应用模型的示例：

def predict_class(model, sentence, min_len=5):
    model.eval()
    tokenized = [tok.text for tok in nlp.tokenizer(sentence)]
    if len(tokenized) < min_len:
        tokenized += ['<pad>'] * (min_len - len(tokenized))
    indexed = [TEXT.vocab.stoi[t] for t in tokenized]
    tensor = torch.LongTensor(indexed).to(device)
    tensor = tensor.unsqueeze(1)
    pred = torch.sigmoid(model(tensor))
    return pred.item()

在这个示例中，我们先对输入的句子进行分词，并将每个单词转换为整数。然后，我们使用predict_class函数传递整个句子，并使用torch.sigmoid将输出转换为概率。

Tags: ,