异常检测在工业中的应用

# 工业异常检测架构概览
import torch
from torch import nn

class IndustrialAnomalyDetector(nn.Module):
    """工业异常检测基础架构"""
    def __init__(self, backbone: str = "resnet18", latent_dim: int = 128):
        super().__init__()
        self.encoder = self._build_backbone(backbone)
        self.bottleneck = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(512, latent_dim),
            nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 512),
            nn.Unflatten(1, (512, 1, 1)),
            nn.ConvTranspose2d(512, 3, kernel_size=4, stride=2, padding=1)
        )

    def _build_backbone(self, name):
        from torchvision.models import resnet18, ResNet18_Weights
        model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
        return nn.Sequential(*list(model.children())[:-2])

    def forward(self, x):
        features = self.encoder(x)
        latent = self.bottleneck(features)
        reconstruction = self.decoder(latent)
        return reconstruction, latent

# 异常分数计算模块
import numpy as np
from sklearn.neighbors import NearestNeighbors

class AnomalyScorer:
    """计算异常分数并生成热力图"""
    def __init__(self, n_neighbors: int = 5):
        self.nn = NearestNeighbors(n_neighbors=n_neighbors, metric="euclidean")
        self.memory_bank = None

    def fit(self, features: np.ndarray):
        self.nn.fit(features)
        self.memory_bank = features
        print(f"Memory bank built with {len(features)} normal samples")

    def score(self, query: np.ndarray) -> np.ndarray:
        distances, _ = self.nn.kneighbors(query)
        anomaly_scores = distances.mean(axis=1)
        normalized = (anomaly_scores - anomaly_scores.min()) / (anomaly_scores.max() - anomaly_scores.min() + 1e-8)
        return normalized

    def generate_heatmap(self, scores: np.ndarray, shape: tuple) -> np.ndarray:
        heatmap = scores.reshape(shape)
        from scipy.ndimage import gaussian_filter
        return gaussian_filter(heatmap, sigma=3)

# PatchCore 核心实现
import torch
from torch.nn import functional as F

class PatchCoreMemoryBank:
    """PatchCore 记忆库构建与推理"""
    def __init__(self, coreset_ratio: float = 0.1):
        self.coreset_ratio = coreset_ratio
        self.memory = None

    def build_memory(self, features: torch.Tensor):
        B, C, H, W = features.shape
        patches = features.permute(0, 2, 3, 1).reshape(-1, C)
        # 随机子采样构建核心集
        n_samples = int(len(patches) * self.coreset_ratio)
        indices = torch.randperm(len(patches))[:n_samples]
        self.memory = patches[indices]
        print(f"Memory bank: {self.memory.shape}")

    def compute_anomaly_map(self, query_features: torch.Tensor) -> torch.Tensor:
        B, C, H, W = query_features.shape
        patches = query_features.permute(0, 2, 3, 1).reshape(-1, C)
        # 归一化
        patches = F.normalize(patches, dim=1)
        memory = F.normalize(self.memory, dim=1)
        # 计算最近邻距离
        dists = torch.cdist(patches, memory)
        min_dists = dists.min(dim=1).values
        anomaly_map = min_dists.reshape(B, H, W)
        return anomaly_map

# MVTec AD 数据加载与预处理
import os
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T

class MVTecDataset(Dataset):
    def __init__(self, root: str, class_name: str, is_train: bool = True):
        self.root = root
        self.class_name = class_name
        base_dir = os.path.join(root, class_name, "train" if is_train else "test")
        self.image_paths = []
        self.mask_paths = []
        for subdir in os.listdir(base_dir):
            for fname in os.listdir(os.path.join(base_dir, subdir)):
                self.image_paths.append(os.path.join(base_dir, subdir, fname))
                mask_path = os.path.join(root, class_name, "ground_truth", subdir, fname.replace(".png", "_mask.png"))
                self.mask_paths.append(mask_path if os.path.exists(mask_path) else None)
        self.transform = T.Compose([
            T.Resize((256, 256)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img = Image.open(self.image_paths[idx]).convert("RGB")
        img = self.transform(img)
        return img

# LSTM 自编码器时序异常检测
import torch
import torch.nn as nn

class LSTMAutoEncoder(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 64, num_layers: int = 2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.encoder = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.decoder = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
        self.output_layer = nn.Linear(hidden_dim, input_dim)

    def forward(self, x):
        batch_size = x.size(0)
        h0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim, device=x.device)
        c0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim, device=x.device)
        _, (h_n, c_n) = self.encoder(x, (h0, c0))
        decoded, _ = self.decoder(h_n[-1].unsqueeze(1).repeat(1, x.size(1), 1).contiguous(), (h_n, c_n))
        reconstructed = self.output_layer(decoded)
        return reconstructed

    def anomaly_score(self, x: torch.Tensor) -> torch.Tensor:
        recon = self.forward(x)
        mse = ((x - recon) ** 2).mean(dim=-1)
        return mse

# 多变量时序异常检测 - OmniAnomaly 简化版
import torch
import torch.nn as nn
import torch.nn.functional as F

class OmniAnomaly(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int = 100, latent_dim: int = 20):
        super().__init__()
        self.encoder_rnn = nn.GRU(input_dim, hidden_dim, batch_first=True)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_var = nn.Linear(hidden_dim, latent_dim)
        self.decoder_rnn = nn.GRU(latent_dim, hidden_dim, batch_first=True)
        self.fc_out = nn.Linear(hidden_dim, input_dim)

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        enc_out, _ = self.encoder_rnn(x)
        mu = self.fc_mu(enc_out)
        logvar = self.fc_var(enc_out)
        z = self.reparameterize(mu, logvar)
        dec_out, _ = self.decoder_rnn(z)
        recon = self.fc_out(dec_out)
        # VAE loss
        recon_loss = F.mse_loss(recon, x, reduction="none").sum(dim=-1)
        kl_loss = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(dim=-1)
        return recon, recon_loss.mean(), kl_loss.mean()

# RUL 预测模型 - CNN-LSTM 混合架构
import torch
import torch.nn as nn

class RULPredictor(nn.Module):
    def __init__(self, n_sensors: int, seq_len: int = 30, hidden: int = 64):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv1d(n_sensors, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv1d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool1d(2)
        )
        self.lstm = nn.LSTM(64, hidden, num_layers=2, batch_first=True, dropout=0.2)
        self.fc = nn.Sequential(
            nn.Linear(hidden, 32),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(32, 1)
        )

    def forward(self, x):
        cnn_out = self.cnn(x.permute(0, 2, 1))
        lstm_in = cnn_out.permute(0, 2, 1)
        lstm_out, _ = self.lstm(lstm_in)
        last_hidden = lstm_out[:, -1, :]
        rul = self.fc(last_hidden)
        return rul.squeeze(-1)

    def predict_with_uncertainty(self, x, n_samples: int = 50):
        predictions = torch.stack([self(x) for _ in range(n_samples)])
        mean = predictions.mean(dim=0)
        std = predictions.std(dim=0)
        return mean, std

# 退化阶段划分 - 隐马尔可夫模型
import numpy as np
from hmmlearn import hmm

class DegradationHMM:
    def __init__(self, n_states: int = 4):
        self.n_states = n_states
        self.model = hmm.GaussianHMM(
            n_components=n_states,
            covariance_type="full",
            n_iter=100
        )
        self.state_labels = ["健康", "轻度退化", "中度退化", "严重退化"]

    def fit(self, sequences: np.ndarray, lengths: np.ndarray):
        self.model.fit(sequences, lengths=lengths)
        print(f"HMM trained with {self.n_states} degradation states")

    def decode_state(self, observation: np.ndarray) -> tuple:
        states = self.model.predict(observation)
        probs = self.model.predict_proba(observation)
        current_state = states[-1]
        confidence = probs[-1, current_state]
        return self.state_labels[current_state], confidence

    def transition_probability(self) -> np.ndarray:
        return self.model.transmat_

# 跨模态注意力融合模块
import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossModalAttention(nn.Module):
    def __init__(self, visual_dim: int, temporal_dim: int, fusion_dim: int = 128):
        super().__init__()
        self.visual_proj = nn.Linear(visual_dim, fusion_dim)
        self.temporal_proj = nn.Linear(temporal_dim, fusion_dim)
        self.cross_attn = nn.MultiheadAttention(fusion_dim, num_heads=4, batch_first=True)
        self.fusion = nn.Sequential(
            nn.Linear(fusion_dim * 2, fusion_dim),
            nn.LayerNorm(fusion_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(fusion_dim, 64)
        )

    def forward(self, visual_feat: torch.Tensor, temporal_feat: torch.Tensor) -> torch.Tensor:
        v = self.visual_proj(visual_feat).unsqueeze(1)
        t = self.temporal_proj(temporal_feat).unsqueeze(1)
        # 交叉注意力
        v_attended, _ = self.cross_attn(v, t, t)
        t_attended, _ = self.cross_attn(t, v, v)
        fused = torch.cat([v_attended.squeeze(1), t_attended.squeeze(1)], dim=-1)
        return self.fusion(fused)

# 多模态工业数据集构建
import torch
from torch.utils.data import Dataset
import numpy as np

class MultiModalIndustrialDataset(Dataset):
    def __init__(self, image_dir: str, sensor_file: str, acoustic_dir: str):
        import pandas as pd
        self.sensor_data = pd.read_csv(sensor_file).values.astype(np.float32)
        self.image_paths = [f"{image_dir}/img_{i}.png" for i in range(len(self.sensor_data))]
        self.acoustic_paths = [f"{acoustic_dir}/audio_{i}.npy" for i in range(len(self.sensor_data))]
        self.labels = np.zeros(len(self.sensor_data))

    def __len__(self):
        return len(self.sensor_data)

    def __getitem__(self, idx):
        from PIL import Image
        import torchvision.transforms as T
        img = Image.open(self.image_paths[idx]).convert("RGB")
        img = T.Compose([T.Resize(224), T.ToTensor()])(img)
        sensor = torch.tensor(self.sensor_data[idx])
        acoustic = torch.tensor(np.load(self.acoustic_paths[idx]))
        return {"image": img, "sensor": sensor, "acoustic": acoustic}

# 模型量化与 TensorRT 部署
import torch
import torch.nn as nn

def quantize_model(model: nn.Module, calib_loader: torch.utils.data.DataLoader):
    model.eval()
    model.qconfig = torch.quantization.get_default_qconfig("fbgemm")
    torch.quantization.prepare(model, inplace=True)
    with torch.no_grad():
        for batch in calib_loader:
            model(batch)
    torch.quantization.convert(model, inplace=True)
    original_size = sum(p.numel() * p.element_size() for p in model.parameters())
    print(f"Original model size: {original_size / 1024:.1f} KB")
    return model

def export_to_onnx(model: nn.Module, dummy_input: torch.Tensor, path: str):
    torch.onnx.export(model, dummy_input, path,
                      export_params=True,
                      opset_version=14,
                      input_names=["input"],
                      output_names=["output"],
                      dynamic_axes={"input": {0: "batch"}})
    print(f"Model exported to {path}")

# 边缘设备推理服务 - FastAPI
from fastapi import FastAPI
import torch
import numpy as np
from pydantic import BaseModel
import base64
from io import BytesIO
from PIL import Image
import torchvision.transforms as T

app = FastAPI(title="工业异常检测 API")
model = torch.jit.load("anomaly_detector.pt")
model.eval()

class InferenceRequest(BaseModel):
    image_b64: str
    threshold: float = 0.5

@app.post("/detect")
def detect_anomaly(req: InferenceRequest):
    img_bytes = base64.b64decode(req.image_b64)
    img = Image.open(BytesIO(img_bytes)).convert("RGB")
    transform = T.Compose([
        T.Resize((256, 256)),
        T.ToTensor(),
        T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    tensor = transform(img).unsqueeze(0)
    with torch.no_grad():
        score = model(tensor).item()
    is_anomaly = score > req.threshold
    return {"anomaly_score": round(score, 4), "is_anomaly": is_anomaly}

# 完整的 PatchCore 训练与评估流程
import torch
from torchvision.models import resnet18, ResNet18_Weights
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
import numpy as np

class CompletePatchCore:
    def __init__(self, backbone: str = "resnet18", coreset_ratio: float = 0.01):
        self.feature_extractor = self._build_extractor()
        self.memory = None
        self.coreset_ratio = coreset_ratio

    def _build_extractor(self):
        model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
        layers = list(model.children())
        return torch.nn.Sequential(*layers[:-2]).eval()

    def extract_features(self, images: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            return self.feature_extractor(images)

    def train(self, normal_images: torch.Tensor, batch_size: int = 32):
        all_features = []
        for i in range(0, len(normal_images), batch_size):
            batch = normal_images[i:i+batch_size]
            feat = self.extract_features(batch)
            all_features.append(feat.cpu())
        all_features = torch.cat(all_features, dim=0)
        # 构建记忆库
        n_patches = all_features.shape[0] * all_features.shape[2] * all_features.shape[3]
        flat = all_features.permute(0, 2, 3, 1).reshape(-1, all_features.shape[1])
        n_select = int(n_patches * self.coreset_ratio)
        indices = torch.randperm(len(flat))[:n_select]
        self.memory = flat[indices]
        print(f"Memory bank: {self.memory.shape}")

    def evaluate(self, test_images: torch.Tensor, labels: torch.Tensor) -> dict:
        scores = []
        for img in test_images:
            feat = self.extract_features(img.unsqueeze(0))
            patches = feat.permute(0, 2, 3, 1).reshape(-1, feat.shape[1])
            patches = F.normalize(patches, dim=1)
            mem = F.normalize(self.memory, dim=1)
            dists = torch.cdist(patches, mem)
            scores.append(dists.min(dim=1).values.max().item())
        auc = roc_auc_score(labels.numpy(), scores)
        return {"image_auc": auc, "n_test": len(test_images)}

# MLOps 流水线 - 模型版本管理与监控
import mlflow
import json
from datetime import datetime

class ModelRegistry:
    def __init__(self, experiment_name: str = "industrial_anomaly_detection"):
        mlflow.set_experiment(experiment_name)
        self.experiment = mlflow.get_experiment_by_name(experiment_name)

    def log_model(self, model: object, metrics: dict, tags: dict):
        with mlflow.start_run() as run:
            mlflow.log_params(tags)
            mlflow.log_metrics(metrics)
            mlflow.log_artifact("config.json")
            model_uri = mlflow.pytorch.log_model(model, "model")
            run_id = run.info.run_id
            print(f"Model logged: run_id={run_id}")
            return run_id

    def promote_to_production(self, run_id: str):
        client = mlflow.tracking.MlflowClient()
        client.transition_model_version_stage(
            name="PatchCoreDetector",
            version=1,
            stage="Production"
        )
        print(f"Model {run_id} promoted to Production")

    def monitor_drift(self, current_data: np.ndarray, reference_data: np.ndarray) -> dict:
        from scipy.stats import ks_2statistic
        drift_scores = {}
        for i in range(current_data.shape[1]):
            stat, p_value = ks_2statistic(reference_data[:, i], current_data[:, i])
            drift_scores[f"feature_{i}"] = {"ks_stat": stat, "p_value": p_value}
        return drift_scores