4. 基于Transformer的故障预测

任务说明

本案例使用Transformer模型对设备运行的时间序列数据进行故障预测。输入为传感器采集的多变量时间序列（如振动、温度、压力），输出为未来一段时间内发生故障的概率。为体现对比价值，同时实现一个基于LSTM的基线模型，在同一数据集上比较两者的预测性能。

数据准备

由于本案例聚焦于模型实现与对比方法，采用公开可复现的模拟数据生成方式：构造一个包含正常段与退化段的二维时间序列（特征1为缓慢漂移变量，特征2为叠加噪声的周期变量），故障标签在退化段起始后置为1。数据按时间窗口切分为样本，每个样本用过去24个时间步预测未来8个时间步内是否发生故障。

Transformer模型结构

本案例实现一个轻量级Transformer编码器用于时序分类：输入经过线性嵌入后加入位置编码，经过2层多头自注意力（4头，d_model=32）和前馈网络，取序列最后一个时间步的表示经全连接层输出二分类概率。该结构不包含因果掩码，因为使用的是编码器而非解码器。

代码实现

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import roc_auc_score, f1_score

# ---------- 1. 数据生成 ----------
def generate_data(n_samples=2000, seq_len=24, n_features=2, seed=42):
    np.random.seed(seed)
    X, y = [], []
    for _ in range(n_samples):
        # 随机起始相位
        t = np.arange(seq_len)
        # 特征1：缓慢漂移 + 噪声
        f1 = np.cumsum(np.random.randn(seq_len) * 0.02) + 0.5
        # 特征2：周期信号 + 噪声
        f2 = np.sin(2 * np.pi * t / 12 + np.random.rand() * 2 * np.pi) + np.random.randn(seq_len) * 0.1
        data = np.stack([f1, f2], axis=1)
        # 故障标签：若特征1漂移超过阈值则视为退化
        label = 1 if np.max(f1) > 0.8 else 0
        X.append(data)
        y.append(label)
    return np.array(X, dtype=np.float32), np.array(y, dtype=np.float32)

# ---------- 2. Transformer模型 ----------
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=100):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

class TransformerPredictor(nn.Module):
    def __init__(self, n_features, d_model=32, nhead=4, num_layers=2):
        super().__init__()
        self.embed = nn.Linear(n_features, d_model)
        self.pos_enc = PositionalEncoding(d_model)
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
                                                   dim_feedforward=64, batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.fc = nn.Linear(d_model, 1)

    def forward(self, x):
        x = self.embed(x)          # (B, L, d_model)
        x = self.pos_enc(x)
        x = self.encoder(x)
        x = x[:, -1, :]            # 取最后时间步
        return torch.sigmoid(self.fc(x)).squeeze(-1)

# ---------- 3. LSTM基线模型 ----------
class LSTMPredictor(nn.Module):
    def __init__(self, n_features, hidden_size=32, num_layers=1):
        super().__init__()
        self.lstm = nn.LSTM(n_features, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        out, _ = self.lstm(x)      # (B, L, hidden)
        out = out[:, -1, :]
        return torch.sigmoid(self.fc(out)).squeeze(-1)

# ---------- 4. 训练与评估 ----------
def train_eval(model, train_loader, test_loader, epochs=10, lr=1e-3):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.BCELoss()
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_loader:
            optimizer.zero_grad()
            loss = criterion(model(xb), yb)
            loss.backward()
            optimizer.step()
    # 评估
    model.eval()
    preds, labels = [], []
    with torch.no_grad():
        for xb, yb in test_loader:
            preds.append(model(xb).numpy())
            labels.append(yb.numpy())
    preds = np.concatenate(preds)
    labels = np.concatenate(labels)
    auc = roc_auc_score(labels, preds)
    f1 = f1_score(labels, (preds > 0.5).astype(int))
    return auc, f1

# ---------- 5. 主流程 ----------
if __name__ == '__main__':
    # 生成数据并划分
    X, y = generate_data()
    split = int(0.8 * len(X))
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    train_ds = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
    test_ds = TensorDataset(torch.from_numpy(X_test), torch.from_numpy(y_test))
    train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_ds, batch_size=32)

    # Transformer
    tf_model = TransformerPredictor(n_features=2)
    tf_auc, tf_f1 = train_eval(tf_model, train_loader, test_loader)

    # LSTM基线
    lstm_model = LSTMPredictor(n_features=2)
    lstm_auc, lstm_f1 = train_eval(lstm_model, train_loader, test_loader)

    # 输出对比结果
    print(f"Transformer  AUC: {tf_auc:.4f}  F1: {tf_f1:.4f}")
    print(f"LSTM        AUC: {lstm_auc:.4f}  F1: {lstm_f1:.4f}")

运行与验证

运行环境需安装torch、numpy、scikit-learn。执行python transformer_fault_prediction.py后，程序将自动完成数据生成、模型训练与评估，并在终端打印两个模型的AUC与F1分数。判定标准：程序无异常退出，且输出两组0到1之间的指标数值。Transformer的AUC通常高于LSTM基线，但具体数值取决于随机数据分布，本案例关注的是对比流程的完整性与可复现性，而非特定性能阈值。