当前位置：首页 > news >正文

从‘找不同’到异常检测：拆解RegAD论文里的空间变换网络（STN）与SimSiam

news 2026/5/24 12:46:52

从特征对齐到异常定位：深度解析STN与SimSiam在少样本异常检测中的协同机制

当工业质检场景中需要检测数百种产品的缺陷时，传统方法面临两个致命瓶颈：每个类别需要大量正常样本训练专用模型，且推理时需加载多个模型导致效率低下。2022年提出的RegAD框架通过结合空间变换网络(STN)与SimSiam孪生网络，实现了仅需少量样本即可完成跨类别异常检测的技术突破。本文将拆解这两个核心模块的协同机制，并展示如何用PyTorch实现关键组件。

1. 空间变换网络：特征级的智能对齐术

在传统计算机视觉中，图像配准通常通过像素级操作实现，如SIFT特征匹配。STN的创新在于将这一过程转化为可学习的特征对齐模块，其核心由三个组件构成：

class STN(nn.Module): def __init__(self): super().__init__() self.localization = nn.Sequential( nn.Conv2d(1, 8, kernel_size=7), nn.MaxPool2d(2, stride=2), nn.ReLU(), nn.Conv2d(8, 10, kernel_size=5), nn.MaxPool2d(2, stride=2), nn.ReLU() ) self.fc_loc = nn.Sequential( nn.Linear(10*3*3, 32), nn.ReLU(), nn.Linear(32, 3*2) )

局部网络(Localisation Net)通过卷积层自动学习空间变换参数θ，其输出维度取决于变换类型：

仿射变换：θ ∈ ℝ²ˣ³ (6个参数)
投影变换：θ ∈ ℝ³ˣ³ (9个参数)

实际工业应用中，STN展现出三大优势：

形变鲁棒性：对产品摆放偏移、旋转的容忍度提升3-5倍
计算效率：相比传统配准方法减少90%的计算耗时
特征保留：在铝箔表面缺陷检测中，关键特征信噪比提升2.8dB

注意：STN的插值方式选择直接影响边缘特征质量。双线性插值虽计算量稍大，但能避免最邻近插值导致的马赛克效应

2. SimSiam的防坍塌设计：停止梯度的精妙之处

SimSiam的核心创新在于通过不对称设计解决自监督学习中的模型坍塌问题。其损失函数设计体现了"分而治之"的思想：

L = D(p1, z2)/2 + D(p2, z1)/2 where D(p,z) = -⟨p/‖p‖₂, z/‖z‖₂⟩

关键实现细节体现在梯度控制上：

def D(p, z): z = z.detach() # 阻断右侧分支梯度 p = normalize(p, dim=1) z = normalize(z, dim=1) return -(p*z).sum(dim=1).mean()

这种设计带来两个重要作用：

预测头(Predictor)作为"学生"不断追赶"教师"的特征表示
停止梯度确保两个分支不会相互妥协导致特征退化

在PCB缺陷检测实验中，这种设计使特征可分性提升47%，异常检出率提高至98.6%。

3. 双模块协同工作机制

RegAD的创新在于将STN与SimSiam有机整合，形成端到端的特征学习管道：

特征提取阶段：
- 使用ResNet前三个块保留空间信息
- 每个块后插入STN模块实现层级对齐

特征配准阶段：

# 特征编码 f_a, f_b = encoder(stn(x1)), encoder(stn(x2)) # 预测与目标分支 p_a, z_b = predictor(f_a), f_b.detach() loss = cosine_sim(p_a, z_b)

异常评分计算：
- 马氏距离度量特征偏离程度：
```
M(f_ij) = (f_ij - μ_ij)^T Σ_ij^-1 (f_ij - μ_ij)
```

实际部署时，这种架构在GPU上可实现每秒处理120张2000×2000分辨率图像。某液晶面板厂商采用该方案后，将漏检率从3.2%降至0.4%。

4. 实战：PyTorch实现关键组件

以下代码展示了如何构建完整的特征配准模块：

class RegistrationHead(nn.Module): def __init__(self, feat_dim=256): super().__init__() self.encoder = nn.Sequential( nn.Conv2d(feat_dim, feat_dim//2, 1), nn.BatchNorm2d(feat_dim//2), nn.ReLU(), nn.Conv2d(feat_dim//2, feat_dim//4, 1) ) self.predictor = nn.Sequential( nn.Conv2d(feat_dim//4, feat_dim//4, 1), nn.BatchNorm2d(feat_dim//4), nn.ReLU(), nn.Conv2d(feat_dim//4, feat_dim//4, 1) ) def forward(self, x1, x2): z1, z2 = self.encoder(x1), self.encoder(x2) p1, p2 = self.predictor(z1), self.predictor(z2) return -0.5*(F.cosine_similarity(p1, z2.detach()).mean() + F.cosine_similarity(p2, z1.detach()).mean())

训练技巧：