Transformer训练例子

import torch

import math

import torch.nn as nn

import torch.optim as optim

from torch.utils.data import Dataset, DataLoader

 # 定义模型

class Transformer(nn.Module):

    def __init__(self, input_size, output_size, hidden_size=512, num_layers=6, num_heads=8):

        super().__init__()

        self.embedding = nn.Embedding(input_size, hidden_size)

        self.pos_encoder = PositionalEncoding(hidden_size, 0.1)

        encoder_layer = nn.TransformerEncoderLayer(hidden_size, num_heads)

        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)

        self.decoder = nn.Linear(hidden_size, output_size)

    def forward(self, src):

        src = self.embedding(src)

        src = self.pos_encoder(src)

        output = self.transformer_encoder(src)

        output = self.decoder(output)

        return output

 # 定义数据集

class MyDataset(Dataset):

    def __init__(self, data, labels):

        self.data = data

        self.labels = labels

    def __len__(self):

        return len(self.data)

    def __getitem__(self, index):

        return self.data[index], self.labels[index]

 # 定义位置编码层

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, dropout=0.1, max_len=5000):

        super().__init__()

        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)

        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)

        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0).transpose(0, 1)

        self.register_buffer('pe', pe)

    def forward(self, x):

        x = x + self.pe[:x.size(0), :]

        return self.dropout(x)

# 进行训练

def train(model, train_loader, criterion, optimizer, scheduler=None, num_epochs=10):

    for epoch in range(num_epochs):

        for data, label in train_loader:

            data = data.to(device)

            label = label.to(device)

            optimizer.zero_grad()

            output = model(data)

            loss = criterion(output, label)

            loss.backward()

            optimizer.step()

        if scheduler is not None:

            scheduler.step()

        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 示例数据

data = [[1, 2, 3, 4, 5], [5, 4, 3, 2, 1], [2, 3, 1, 5, 4], [4, 2, 1, 3, 5]]

labels = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]

 # 转换为Tensor类型

data = torch.LongTensor(data)

labels = torch.LongTensor(labels)

 # 定义超参数

input_size = 6

output_size = 4

hidden_size = 512

num_layers = 6

num_heads = 8

batch_size = 2

num_epochs = 100

learning_rate = 0.01

device='cpu'

 # 创建模型和数据集

model = Transformer(input_size, output_size, hidden_size, num_layers, num_heads).to(device)

dataset = MyDataset(data, labels)

 # 定义数据加载器和损失函数

train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

criterion = nn.CrossEntropyLoss()

 # 定义优化器和学习率调度器

optimizer = optim.Adam(model.parameters(), lr=learning_rate)

scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)

 # 开始训练

train(model, train_loader, criterion, optimizer, scheduler, num_epochs)