Transformer 例子

据说很好用，先写一个例子看看：

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 构造简单的时间序列数据集
def generate_time_series():
    time = np.arange(0, 100, 0.1)
    amplitude = np.sin(time)
    return amplitude

# 将时间序列数据转换为序列样本
def create_sequences(data, seq_length):
    sequences = []
    for i in range(len(data) - seq_length):
        seq = data[i:i + seq_length]
        label = data[i + seq_length]
        sequences.append((seq, label))
    return sequences

# 定义 Transformer 模型
class TransformerModel(nn.Module):
    def __init__(self, input_size, output_size, num_layers, heads, hidden_size):
        super(TransformerModel, self).__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=input_size, nhead=heads, dim_feedforward=hidden_size)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        self.decoder = nn.Linear(input_size, output_size)

    def forward(self, src):
        src = src.unsqueeze(0)  # 添加批次维度
        output = self.transformer_encoder(src)
        output = self.decoder(output[-1])  # 取最后一个时间步的输出
        return output

# 定义训练函数
def train(model, criterion, optimizer, epochs, train_loader):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for inputs, labels in train_loader:
            optimizer.zero_grad()
            outputs = model(inputs.float())
            loss = criterion(outputs.squeeze(), labels.float())
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss}")

# 设置随机种子，保证实验的可复现性
torch.manual_seed(0)
np.random.seed(0)

# 生成时间序列数据并创建序列样本
data = generate_time_series()
seq_length = 10
sequences = create_sequences(data, seq_length)

# 将序列样本划分为训练集和测试集
train_size = int(len(sequences) * 0.8)
train_data = sequences[:train_size]
test_data = sequences[train_size:]

# 准备训练数据加载器
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)

# 定义模型和优化器
model = TransformerModel(input_size=seq_length, output_size=1, num_layers=2, heads=2, hidden_size=128)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练模型
train(model, criterion, optimizer, epochs=20, train_loader=train_loader)

# 测试模型
model.eval()
# 测试所有测试数据
test_inputs = torch.tensor([test_data[i][0] for i in range(len(test_data))]).float()
with torch.no_grad():
    predicted = model(test_inputs).squeeze().numpy()

test_labels = [test_data[i][1] for i in range(len(test_data))]
# 可视化预测结果
plt.plot(predicted, label='predicted')
plt.plot(test_labels, label='test')
plt.legend()
plt.show()