PyTorch中使用Transformer对一维序列进行分类的源代码

在PyTorch中使用Transformer对一维序列进行分类是一种常见做法，尤其是在处理时间序列数据、自然语言处理等领域。Transformer模型因其并行化处理能力和自注意力机制而成为许多任务的首选模型。

下面是一个使用PyTorch实现Transformer对一维序列进行分类的完整示例代码，包括数据预处理、模型定义、训练和评估等部分。

1. 准备工作

首先，我们需要导入所需的库，并定义一些基本的参数。

1import torch
2import torch.nn as nn
3import torch.optim as optim
4from torch.utils.data import Dataset, DataLoader
5from sklearn.model_selection import train_test_split
6from sklearn.preprocessing import StandardScaler
7import numpy as np
8
9# 定义超参数
10input_dim = 10  # 序列的维度
11seq_length = 50  # 序列长度
12hidden_dim = 128  # Transformer编码器隐藏层维度
13num_heads = 8  # 多头注意力机制中的头数
14num_layers = 6  # 编码器层数
15dropout = 0.1  # dropout概率
16num_classes = 3  # 分类类别数
17batch_size = 32
18num_epochs = 100
19learning_rate = 0.001

2. 数据预处理

假设我们有一组一维序列数据，我们将对其进行预处理，并将其划分为训练集和测试集。

1# 生成模拟数据
2def generate_data(n_samples, seq_length, input_dim, num_classes):
3    X = np.random.randn(n_samples, seq_length, input_dim)
4    y = np.random.randint(0, num_classes, size=(n_samples,))
5    return X, y
6
7# 生成数据
8X, y = generate_data(1000, seq_length, input_dim, num_classes)
9
10# 数据标准化
11scaler = StandardScaler()
12X = scaler.fit_transform(X.reshape(-1, input_dim)).reshape(X.shape)
13
14# 划分数据集
15X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
16
17# 定义数据集类
18class SequenceDataset(Dataset):
19    def __init__(self, X, y):
20        self.X = X
21        self.y = y
22
23    def __len__(self):
24        return len(self.y)
25
26    def __getitem__(self, idx):
27        return torch.tensor(self.X[idx], dtype=torch.float32), torch.tensor(self.y[idx], dtype=torch.long)
28
29train_dataset = SequenceDataset(X_train, y_train)
30test_dataset = SequenceDataset(X_test, y_test)
31
32train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
33test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

3. 定义Transformer模型

接下来定义一个基于Transformer的模型，该模型包含嵌入层、位置编码、多头自注意力机制和前馈神经网络等组件。

1class PositionalEncoding(nn.Module):
2    def __init__(self, d_model, max_len=5000):
3        super(PositionalEncoding, self).__init__()
4        pe = torch.zeros(max_len, d_model)
5        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
6        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
7        pe[:, 0::2] = torch.sin(position * div_term)
8        pe[:, 1::2] = torch.cos(position * div_term)
9        pe = pe.unsqueeze(0).transpose(0, 1)
10        self.register_buffer('pe', pe)
11
12    def forward(self, x):
13        return x + self.pe[:x.size(0), :]
14
15class TransformerClassifier(nn.Module):
16    def __init__(self, input_dim, hidden_dim, num_heads, num_layers, num_classes, dropout):
17        super(TransformerClassifier, self).__init__()
18        self.embedding = nn.Linear(input_dim, hidden_dim)
19        self.positional_encoding = PositionalEncoding(hidden_dim)
20        encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=num_heads, dropout=dropout)
21        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)
22        self.classifier = nn.Linear(hidden_dim, num_classes)
23        self.dropout = nn.Dropout(dropout)
24
25    def forward(self, src):
26        embedded = self.embedding(src) * np.sqrt(hidden_dim)
27        encoded = self.positional_encoding(embedded)
28        output = self.transformer_encoder(encoded)
29        output = output.mean(dim=0)  # 平均池化
30        output = self.classifier(output)
31        return output
32
33model = TransformerClassifier(input_dim, hidden_dim, num_heads, num_layers, num_classes, dropout)

4. 训练模型

定义损失函数、优化器，并进行模型训练。

1criterion = nn.CrossEntropyLoss()
2optimizer = optim.Adam(model.parameters(), lr=learning_rate)
3
4def train(model, data_loader, criterion, optimizer, device):
5    model.train()
6    total_loss = 0.0
7    for batch_idx, (data, target) in enumerate(data_loader):
8        data, target = data.to(device), target.to(device)
9        optimizer.zero_grad()
10        output = model(data.permute(1, 0, 2))  # 调整数据维度为 (seq_len, batch_size, input_dim)
11        loss = criterion(output, target)
12        loss.backward()
13        optimizer.step()
14        total_loss += loss.item()
15    return total_loss / (batch_idx + 1)
16
17def evaluate(model, data_loader, criterion, device):
18    model.eval()
19    total_loss = 0.0
20    correct = 0
21    with torch.no_grad():
22        for data, target in data_loader:
23            data, target = data.to(device), target.to(device)
24            output = model(data.permute(1, 0, 2))
25            loss = criterion(output, target)
26            total_loss += loss.item()
27            pred = output.argmax(dim=1, keepdim=True)
28            correct += pred.eq(target.view_as(pred)).sum().item()
29    accuracy = correct / len(data_loader.dataset)
30    return total_loss / len(data_loader), accuracy
31
32device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
33model.to(device)
34
35for epoch in range(num_epochs):
36    train_loss = train(model, train_loader, criterion, optimizer, device)
37    test_loss, test_acc = evaluate(model, test_loader, criterion, device)
38    print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}")

5. 总结

以上代码实现了一个使用Transformer对一维序列进行分类的完整流程，包括数据预处理、模型定义、训练和评估。该模型适用于处理时间序列数据或其他一维序列数据的分类任务。通过调整超参数和网络结构，可以进一步优化模型的性能。