Dinomaly2：面向全谱系无监督异常检测的统一框架

# ============================================================================
# One Dinomaly2 Detect Them All: 统一全谱无监督异常检测框架
# 详细伪代码 - 包含核心算法、训练和推理完整流程
# ============================================================================

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Tuple, List, Dict

# ========================= 核心模块定义 =========================

class MultiScaleFeatureExtractor(nn.Module):
    """多尺度特征提取器"""
    def __init__(self, backbone: nn.Module, scales: List[int] = [1, 2, 4]):
        super().__init__()
        self.backbone = backbone  # 骨干网络（如Vision Transformer）
        self.scales = scales
        
    def forward(self, x: torch.Tensor) -> Dict[int, torch.Tensor]:
        """
        输入: 
            x: 输入图像 [B, C, H, W]
        输出:
            多尺度特征字典 {scale: feature}
        """
        features = {}
        base_feature = self.backbone(x)  # 提取基础特征
        
        for scale in self.scales:
            # 通过池化或卷积下采样获取多尺度特征
            if scale == 1:
                features[scale] = base_feature
            else:
                h, w = base_feature.shape[2:]
                pooled = F.adaptive_avg_pool2d(base_feature, (h//scale, w//scale))
                features[scale] = pooled
        return features

class LearnableMemoryBank(nn.Module):
    """可学习记忆库模块"""
    def __init__(self, feature_dim: int, memory_size: int = 1000):
        super().__init__()
        self.feature_dim = feature_dim
        self.memory_size = memory_size
        
        # 可训练的记忆向量
        self.memory = nn.Parameter(torch.randn(memory_size, feature_dim))
        # 记忆项的使用频率（用于在线更新）
        self.usage_counter = torch.zeros(memory_size)
        
    def query(self, features: torch.Tensor, k: int = 5) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        查询记忆库 - 获取最相关的记忆项
        
        输入:
            features: 查询特征 [B, D, H, W]
            k: 最近邻数量
        输出:
            reconstructed_features: 重构特征 [B, D, H, W]
            residual: 残差特征 [B, D, H, W]
        """
        B, D, H, W = features.shape
        
        # 展平特征以便查询
        flat_features = features.permute(0, 2, 3, 1).reshape(-1, D)  # [B*H*W, D]
        
        # 计算与所有记忆项的相似度
        similarities = F.cosine_similarity(
            flat_features.unsqueeze(1),  # [B*H*W, 1, D]
            self.memory.unsqueeze(0),    # [1, M, D]
            dim=2
        )  # [B*H*W, M]
        
        # 获取top-k最相似的记忆项
        topk_sim, topk_indices = torch.topk(similarities, k=k, dim=1)  # [B*H*W, k]
        
        # 计算注意力权重（softmax归一化）
        attention_weights = F.softmax(topk_sim / 0.1, dim=1)  # 温度系数0.1
        
        # 加权组合记忆项
        selected_memory = self.memory[topk_indices]  # [B*H*W, k, D]
        reconstructed_flat = torch.sum(
            attention_weights.unsqueeze(-1) * selected_memory, 
            dim=1
        )  # [B*H*W, D]
        
        # 更新使用计数
        with torch.no_grad():
            unique_indices = torch.unique(topk_indices)
            self.usage_counter[unique_indices] += 1
        
        # 恢复空间维度
        reconstructed_features = reconstructed_flat.reshape(B, H, W, D).permute(0, 3, 1, 2)
        residual = features - reconstructed_features
        
        return reconstructed_features, residual, attention_weights

class MultiHeadReconstructionDecoder(nn.Module):
    """多头重构解码器（用于像素级异常检测）"""
    def __init__(self, feature_dims: List[int], output_channels: int = 3):
        super().__init__()
        self.decoders = nn.ModuleList([
            nn.Sequential(
                nn.ConvTranspose2d(dim, 128, 3, 2, 1),
                nn.BatchNorm2d(128),
                nn.ReLU(),
                nn.ConvTranspose2d(128, 64, 3, 2, 1),
                nn.BatchNorm2d(64),
                nn.ReLU(),
                nn.Conv2d(64, output_channels, 1)
            ) for dim in feature_dims
        ])
        
    def forward(self, multi_scale_features: Dict[int, torch.Tensor]) -> Dict[int, torch.Tensor]:
        """
        输入: 多尺度特征字典
        输出: 多尺度重构图像字典
        """
        reconstructions = {}
        for scale, decoder in zip(sorted(multi_scale_features.keys()), self.decoders):
            features = multi_scale_features[scale]
            reconstructions[scale] = decoder(features)
        return reconstructions

class DiscriminativeDecoder(nn.Module):
    """判别解码器（用于图像级异常检测）"""
    def __init__(self, feature_dim: int, hidden_dim: int = 256):
        super().__init__()
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(feature_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )
        
    def forward(self, residual_features: torch.Tensor) -> torch.Tensor:
        """
        输入: 残差特征 [B, D, H, W]
        输出: 异常概率 [B, 1]
        """
        return self.classifier(residual_features)

# ========================= 完整框架定义 =========================

class Dinomaly2(nn.Module):
    """统一全谱异常检测框架"""
    def __init__(
        self,
        backbone: nn.Module,
        feature_dim: int = 256,
        memory_size: int = 1000,
        scales: List[int] = [1, 2, 4]
    ):
        super().__init__()
        
        # 特征提取部分
        self.feature_extractor = MultiScaleFeatureExtractor(backbone, scales)
        
        # 记忆库
        self.memory_bank = LearnableMemoryBank(feature_dim, memory_size)
        
        # 重构解码器（多头，对应不同尺度）
        self.reconstruction_decoder = MultiHeadReconstructionDecoder(
            [feature_dim] * len(scales)
        )
        
        # 判别解码器
        self.discriminative_decoder = DiscriminativeDecoder(feature_dim)
        
        # 损失权重
        self.lambda_recon = 1.0
        self.lambda_disc = 0.1
        self.lambda_commit = 0.01
        
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        输入: 
            x: 输入图像 [B, C, H, W]
        输出:
            包含中间结果和最终输出的字典
        """
        # 1. 提取多尺度特征
        multi_scale_features = self.feature_extractor(x)
        
        # 2. 记忆库查询（对每个尺度）
        reconstructed_features = {}
        residuals = {}
        attention_weights = {}
        
        for scale, features in multi_scale_features.items():
            recon_feat, residual, attn = self.memory_bank.query(features)
            reconstructed_features[scale] = recon_feat
            residuals[scale] = residual
            attention_weights[scale] = attn
        
        # 3. 重构图像
        reconstructions = self.reconstruction_decoder(reconstructed_features)
        
        # 4. 判别分数
        # 使用最细尺度（scale=1）的残差进行判别
        disc_scores = self.discriminative_decoder(residuals[1])
        
        return {
            'multi_scale_features': multi_scale_features,
            'reconstructed_features': reconstructed_features,
            'residuals': residuals,
            'attention_weights': attention_weights,
            'reconstructions': reconstructions,
            'disc_scores': disc_scores
        }

# ========================= 训练流程 =========================

def train_dinomaly2(
    model: Dinomaly2,
    train_loader: torch.utils.data.DataLoader,
    num_epochs: int = 100,
    device: str = 'cuda'
):
    """
    训练函数
    
    假设: 训练数据全部为正常样本
    """
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    # 损失函数
    reconstruction_loss = nn.L1Loss()
    discrimination_loss = nn.BCELoss()
    
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        
        for batch_idx, (images, _) in enumerate(train_loader):
            images = images.to(device)
            
            # 前向传播
            outputs = model(images)
            
            # 计算各项损失
            # 1. 重构损失（多尺度）
            recon_loss = 0
            for scale, recon_img in outputs['reconstructions'].items():
                # 调整重构图像大小以匹配输入
                if recon_img.shape != images.shape:
                    recon_img = F.interpolate(recon_img, size=images.shape[2:])
                recon_loss += reconstruction_loss(recon_img, images)
            
            # 2. 判别损失（正常样本应接近0）
            disc_loss = discrimination_loss(
                outputs['disc_scores'], 
                torch.zeros_like(outputs['disc_scores'])
            )
            
            # 3. 记忆库提交损失（鼓励特征被记忆库充分表示）
            commit_loss = 0
            for scale, residual in outputs['residuals'].items():
                commit_loss += torch.mean(torch.norm(residual, p=2, dim=1))
            
            # 总损失
            loss = (
                model.lambda_recon * recon_loss +
                model.lambda_disc * disc_loss +
                model.lambda_commit * commit_loss
            )
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            
            if batch_idx % 100 == 0:
                print(f'Epoch {epoch}, Batch {batch_idx}: Loss = {loss.item():.4f}')
        
        # 定期更新记忆库（移除不常用的记忆项）
        if epoch % 10 == 0:
            update_memory_bank(model.memory_bank)
        
        print(f'Epoch {epoch} completed. Average Loss: {total_loss/len(train_loader):.4f}')

def update_memory_bank(memory_bank: LearnableMemoryBank, usage_threshold: float = 0.1):
    """
    更新记忆库 - 替换不常用的记忆项
    """
    with torch.no_grad():
        usage = memory_bank.usage_counter / memory_bank.usage_counter.sum()
        
        # 找出使用频率低于阈值的记忆项索引
        rarely_used_indices = torch.where(usage < usage_threshold)[0]
        
        if len(rarely_used_indices) > 0:
            # 用随机噪声替换这些记忆项（模拟新记忆的引入）
            new_memory = torch.randn(len(rarely_used_indices), memory_bank.feature_dim)
            memory_bank.memory.data[rarely_used_indices] = new_memory
            
            # 重置这些记忆项的使用计数
            memory_bank.usage_counter[rarely_used_indices] = 0

# ========================= 推理流程 =========================

def inference_dinomaly2(
    model: Dinomaly2, 
    test_images: torch.Tensor,
    anomaly_threshold: float = 0.5
) -> Dict[str, np.ndarray]:
    """
    推理函数 - 生成像素级和图像级异常检测结果
    
    输入:
        model: 训练好的模型
        test_images: 测试图像 [B, C, H, W]
        anomaly_threshold: 异常阈值
    输出:
        包含各种异常分数的字典
    """
    model.eval()
    with torch.no_grad():
        outputs = model(test_images)
        
        # 提取关键结果
        original_images = test_images.cpu().numpy()
        reconstructions = outputs['reconstructions'][1].cpu().numpy()  # 取最细尺度
        disc_scores = outputs['disc_scores'].cpu().numpy()
        
        # 计算像素级异常热图
        B, C, H, W = original_images.shape
        
        # 1. 重构误差热图
        pixel_recon_error = np.abs(original_images - reconstructions)
        pixel_recon_error = np.mean(pixel_recon_error, axis=1)  # [B, H, W]
        
        # 2. 残差特征幅度热图
        residual_features = outputs['residuals'][1].cpu().numpy()  # [B, D, H, W]
        residual_magnitude = np.linalg.norm(residual_features, axis=1)  # [B, H, W]
        
        # 3. 记忆注意力不确定性热图
        attention_weights = outputs['attention_weights'][1].cpu().numpy()  # [B*H*W, k]
        # 计算注意力分布的熵作为不确定性
        attn_entropy = -np.sum(attention_weights * np.log(attention_weights + 1e-10), axis=1)
        attention_uncertainty = attn_entropy.reshape(B, H, W)
        
        # 融合多源异常信号（像素级）
        alpha, beta = 0.5, 0.3  # 融合权重
        pixel_anomaly_map = (
            alpha * normalize(pixel_recon_error) +
            beta * normalize(residual_magnitude) +
            (1 - alpha - beta) * normalize(attention_uncertainty)
        )
        
        # 图像级异常分数融合
        image_anomaly_score = (
            0.6 * disc_scores.squeeze() + 
            0.4 * np.max(pixel_anomaly_map.reshape(B, -1), axis=1)
        )
        
        # 二值化异常检测结果
        pixel_anomaly_binary = (pixel_anomaly_map > anomaly_threshold).astype(np.uint8)
        image_anomaly_binary = (image_anomaly_score > anomaly_threshold).astype(np.uint8)
        
    return {
        'pixel_anomaly_map': pixel_anomaly_map,          # 像素级异常热图 [B, H, W]
        'image_anomaly_score': image_anomaly_score,      # 图像级异常分数 [B]
        'pixel_anomaly_binary': pixel_anomaly_binary,    # 像素级二值化结果 [B, H, W]
        'image_anomaly_binary': image_anomaly_binary,    # 图像级二值化结果 [B]
        'reconstructions': reconstructions,              # 重构图像 [B, C, H, W]
        'attention_weights': attention_weights           # 注意力权重 [B*H*W, k]
    }

def normalize(x: np.ndarray) -> np.ndarray:
    """归一化到[0,1]范围"""
    x_min = x.min(axis=(1, 2), keepdims=True)
    x_max = x.max(axis=(1, 2), keepdims=True)
    return (x - x_min) / (x_max - x_min + 1e-10)

# ========================= 示例使用 =========================

if __name__ == "__main__":
    # 1. 初始化模型
    # 假设我们有一个预训练的Vision Transformer作为骨干
    from transformers import ViTModel
    
    vit = ViTModel.from_pretrained('google/vit-base-patch16-224')
    vit_backbone = vit.embeddings  # 简化示意，实际需要调整
    
    model = Dinomaly2(
        backbone=vit_backbone,
        feature_dim=256,
        memory_size=1000,
        scales=[1, 2, 4]
    )
    
    # 2. 训练模型（假设已有正常样本的数据加载器）
    # train_dinomaly2(model, train_loader, num_epochs=100)
    
    # 3. 推理测试
    # 创建模拟测试数据
    batch_size = 4
    channels = 3
    height, width = 224, 224
    test_images = torch.randn(batch_size, channels, height, width)
    
    # 进行推理
    results = inference_dinomaly2(model, test_images, anomaly_threshold=0.5)
    
    print("推理结果:")
    print(f"- 像素级异常图形状: {results['pixel_anomaly_map'].shape}")
    print(f"- 图像级异常分数: {results['image_anomaly_score']}")
    print(f"- 检测到异常图像数量: {results['image_anomaly_binary'].sum()}")