import torch
import torch.nn as nn
import torch.optim as optim
import math
from torchtext.legacy.data import Field, BucketIterator
from torchtext.legacy import datasets

# 定义模型
class TransformerModel(nn.Module):
    def __init__(self, input_dim, emb_dim, n_heads, hid_dim, n_layers, output_dim, dropout):
        super().__init__()

        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.pos_encoding = PositionalEncoding(emb_dim, dropout)
        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(emb_dim, n_heads, hid_dim, dropout), n_layers
        )
        self.fc = nn.Linear(emb_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        # src: [src_len, batch_size]
        embedded = self.embedding(src) * math.sqrt(self.emb_dim)
        embedded = self.pos_encoding(embedded)
        outputs = self.transformer_encoder(embedded)
        last_output = outputs[-1, :, :]
        last_output = self.dropout(last_output)
        return self.fc(last_output)

1.2、定义了一个PositionalEncoding类，用于实现位置编码；

# 定义位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, emb_dim, dropout, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, emb_dim, 2) * -(math.log(10000.0) / emb_dim))
        pe = torch.zeros(max_len, 1, emb_dim)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

1.3、定义了train和evaluate函数，用于训练和验证模型；

# 定义训练函数
def train(model, iterator, optimizer, criterion):
    epoch_loss = 0
    model.train()

    for batch in iterator:
        optimizer.zero_grad()
        src = batch.text
        trg = batch.label
        output = model(src)
        loss = criterion(output.squeeze(1), trg.float())
        loss.backward()
        optimizer.step()
        if batch.batch_size > 0:  # 判断当前batch中样本的数量是否为0
            epoch_loss += loss.item() * batch.batch_size
    return epoch_loss / len(iterator.dataset)


# 定义验证函数
def evaluate(model, iterator, criterion):
    epoch_loss = 0
    model.eval()

    with torch.no_grad():
        for batch in iterator:
            src = batch.text
            trg = batch.label
            output = model(src)
            loss = criterion(output.squeeze(1), trg.float())
            if batch.batch_size > 0:  # 判断当前batch中样本的数量是否为0
                epoch_loss += loss.item() * batch.batch_size
    return epoch_loss / len(iterator.dataset)

2、主函数

2.1、加载IMDB数据集，并初始化模型和优化器；


if __name__ == '__main__':
    # 设置参数
    BATCH_SIZE = 32
    EMB_DIM = 256
    HID_DIM = 512
    N_LAYERS = 6
    N_HEADS = 8
    DROPOUT = 0.1
    LEARNING_RATE = 0.0005
    N_EPOCHS = 10

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载数据集
    text = Field(tokenize='spacy', tokenizer_language='en_core_web_sm')
    label = Field(dtype=torch.float)
    train_data, test_data = datasets.IMDB.splits(text, label)
    text.build_vocab(train_data, max_size=10000, vectors='glove.6B.100d')
    label.build_vocab(train_data)
    train_iterator, test_iterator = BucketIterator.splits(
        (train_data, test_data), batch_size=BATCH_SIZE, device=device, sort_key=False  )

    # 初始化模型和优化器
    INPUT_DIM = len(text.vocab)
    OUTPUT_DIM = 1
    model = TransformerModel(INPUT_DIM, EMB_DIM, N_HEADS, HID_DIM, N_LAYERS, OUTPUT_DIM, DROPOUT)
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    criterion = nn.BCEWithLogitsLoss()

2.2、训练模型，并输出训练和验证损失。

    # 训练模型
    for epoch in range(N_EPOCHS):
        train_loss = train(model, train_iterator, optimizer, criterion)
        test_loss = evaluate(model, test_iterator, criterion)
        print(f'Epoch: epoch +1:02 | Train Loss: train_loss:.3f | Test Loss: test_loss:.3f')

以上是关于PT之Transformer：基于PyTorch框架利用Transformer算法针对IMDB数据集实现情感分类的应用案例代码解析的主要内容，如果未能解决你的问题，请参考以下文章

PyTorch深度学习项目实战100例—— 基于Transformer实现Twitter文本隐喻二分类 | 第43例

Transformer解读（附pytorch代码）

PyTorch 使用 nn.Transformer 和 TorchText 进行序列到序列的建模

PyTorch-19 使用 nn.Transformer 和 TorchText 进行序列到序列的建模

人工智能(Pytorch)搭建transformer模型,真正跑通transformer模型，深刻了解transformer的架构

基于Transformer的NLP处理管线