Python如何实现基于注意力机制的异常检测？Transformer

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 定义Transformer编码器
class TransformerEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, num_heads):
        super(TransformerEncoder, self).__init__()
        self.embedding = nn.Linear(input_dim, hidden_dim)
        self.transformer_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(hidden_dim, num_heads)
            for _ in range(num_layers)
        ])
        self.linear = nn.Linear(hidden_dim, input_dim)

    def forward(self, x):
        x = self.embedding(x)
        for layer in self.transformer_layers:
            x = layer(x)
        x = self.linear(x)
        return x

# 生成一些随机数据
input_dim = 10
sequence_length = 20
batch_size = 32
num_epochs = 10
learning_rate = 0.001

# 创建模型
model = TransformerEncoder(input_dim, hidden_dim=32, num_layers=2, num_heads=4)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 生成训练数据 (全部是正常数据)
train_data = torch.randn(batch_size, sequence_length, input_dim)

# 训练模型
for epoch in range(num_epochs):
    optimizer.zero_grad()
    output = model(train_data)
    loss = criterion(output, train_data)
    loss.backward()
    optimizer.step()
    print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

# 异常检测
# 生成测试数据 (包含一些异常数据)
test_data = torch.randn(batch_size, sequence_length, input_dim)
# 模拟异常数据 (例如，在某个维度上增加较大的噪声)
test_data[0, :, 0] += 5  # 第一个样本的第一个特征添加噪声

model.eval() # 设置为评估模式
with torch.no_grad():
    output = model(test_data)
    loss = criterion(output, test_data)
    print(f"Test Loss: {loss.item()}")

    # 设置阈值 (例如，基于训练数据的损失分布)
    threshold = 0.1 # 示例阈值，需要根据实际情况调整

    # 异常检测判断
    for i in range(batch_size):
        sample_loss = criterion(output[i:i+1], test_data[i:i+1]).item()
        if sample_loss > threshold:
            print(f"Sample {i+1} is anomalous (Loss: {sample_loss})")
        else:
            print(f"Sample {i+1} is normal (Loss: {sample_loss})")