ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Anomalib 中的 DSR 双 subspace 重投影模型:量化特征、三阶段训练与异常分割实现详解

Anomalib 中的 DSR 双 subspace 重投影模型:量化特征、三阶段训练与异常分割实现详解 Anomalib 中的 DSR 双 subspace 重投影模型量化特征、三阶段训练与异常分割实现详解【免费下载链接】anomalibAn anomaly detection library comprising state-of-the-art algorithms and features such as experiment management, hyper-parameter optimization, and edge inference.项目地址: https://gitcode.com/GitHub_Trending/an/anomalibDSRDual Subspace Re-Projection是 Anomalib 提供的量化特征型表面异常检测算法它以“一编码器 双解码器”的自编码器为核心在 ImageNet 上预训练得到量化码本再分三个阶段训练重建、子空间约束与上采样模块最终输出像素级异常分割图。本文围绕 DSR 的模型文档展开结合 Anomalib 仓库中的 PyTorch 网络定义、Lightning 训练封装、异常生成器与损失函数源码讲清楚 DSR 的架构组成、训练流程、关键超参数latent_anomaly_strength、upsampling_train_ratio以及如何在 Anomalib 中实际训练和推理该模型。1. DSR 的核心思想用量化码本“凭空”制造异常DSR 官方模型文档 docs/source/markdown/guides/reference/models/image/dsr.md 对其描述为一个由自编码器单编码器 双解码器与异常检测模块组成的量化特征算法。其工作流程可概括为预训练量化码本编码器与“通用物体解码器general object decoder”以及码本codebook先在 ImageNet 上完成预训练得到一组量化表征字典特征级缺陷生成利用码本对输入图像的量化表征进行“近分布”采样在特征层面人为制造缺陷从而训练物体专属解码器object-specific decoder和异常检测模块——全程不需要外部缺陷数据集上采样精修最后阶段用模拟的图像级污点smudge训练上采样模块输出更鲁棒的异常图。这种设计的本质是把“正常外观”约束在离散码本张成的子空间内任何无法被码本良好表示的区域都会被判定为异常并通过两路重建图像的差异来定位异常位置。模型类型Model Type为Segmentation即在像素级输出异常分割图而非仅给图像级分数。2. 网络组成从 torch_model.py 看五大模块DSR 的 PyTorch 实现位于 src/anomalib/models/image/dsr/torch_model.py入口类为DsrModel。从源码结构看DsrModel由五个功能块拼装而成模块类名作用离散潜空间模型DiscreteLatentModel将图像编码为 hi/lo 两级量化特征图并用通用解码器重建图像图像重建网络物体专属解码器ImageReconstructionNetwork从量化特征重建图像输入通道为embedding_dim * 2子空间约束模块 ×2SubspaceRestrictionModulehi/lo 各一把受缺陷污染的量化特征“拉回”正常外观子空间并重新量化异常检测模块AnomalyDetectionModule对比两路重建图像输出 2 通道分割 logits上采样模块UpsamplingModule将低分辨率分割图放大到全分辨率2.1 离散潜空间模型与量化码本DiscreteLatentModel见 torch_model.py内部包含EncoderBot/EncoderTop两级编码器。EncoderBot以 RGB 图像为输入先做两次 stride-2 卷积下采样EncoderTop在其输出上继续下采样形成 hi高分辨率特征对应quantized_b与 lo低分辨率特征对应quantized_t两条路径。源码中有注释reminder : top lo, bot hi!说明命名中 top/bot 与 hi/lo 是交叉对应的VectorQuantizerVQ-VAE 风格的量化器默认码本规模num_embeddings4096、向量维度embedding_dim128DsrModel的默认参数。量化过程通过最近邻查表完成torch.argmin one-hot 矩阵乘法并用inputs (quantized - inputs).detach()实现直通估计器straight-through estimator传梯度upsample_t一次ConvTranspose2d把 lo 特征上采样后与 hi 特征拼接共同送入通用解码器DecoderBot得到重建图像。DsrModel.__init__中对discrete_latent_model的全部参数执行parameters.requires_grad False见 torch_model.py#L120-L121这与文档描述一致离散模型阶段第一阶段是 ImageNet 预训练的Anomalib 训练时直接冻结加载不在线训练。2.2 前向推理链路DsrModel.forward在推理/评估分支anomaly_map_to_generate is None中见 torch_model.py#L181-L232在torch.no_grad()下调用离散模型得到通用解码器重建图gen_image与 hi/lo 量化特征对量化特征做副本并送入两个SubspaceRestrictionModule得到“无缺陷”的重投影量化特征再经物体专属解码器重建出obj_spec_image将两路重建图拼接各 3 通道后送入AnomalyDetectionModule一个 U-Netbase_width64输出 2 通道得到低分辨率分割 logitssoftmax 后得到out_mask_sm最后UpsamplingModule接收“物体专属重建图 通用重建图 低分辨率分割图”共 3328 通道对应构造参数in_channels(2 * image_dim) anomaly_map_dim输出全分辨率 2 通道分割图。图像级异常分数则来自对 softmax 后异常类通道做avg_pool2d(kernel21, stride1)再取空间最大值torch.amax见 torch_model.py#L222-L228最终返回InferenceBatch(pred_score..., anomaly_map...)其中anomaly_map取上采样后 softmax 的异常类通道out_mask_sm_up[:, 1, :, :]。3. Lightning 封装三阶段训练与关键超参数对外暴露的训练类是 src/anomalib/models/image/dsr/lightning_model.py 中的Dsr继承AnomalibModulelearning_type为LearningType.ONE_CLASS。它只有两个模型专属超参数参数默认值含义latent_anomaly_strength0.2潜空间中生成异常的强度下界。第二阶段训练时每张图片的实际异常强度在[latent_anomaly_strength, 1.0]内随机采样见 torch_model.py#L238-L243 的torch.rand(...) * (1.0 - strength) strengthupsampling_train_ratio0.7前多少比例的 epoch 用于第二阶段训练之后进入第三阶段上采样模块训练其余构造参数为 Anomalib 通用的pre_processor/post_processor/evaluator/visualizer均可置为True使用默认实现。3.1 预训练权重自动下载Dsr类定义了模块级常量WEIGHTS_DOWNLOAD_INFO见 lightning_model.py#L63-L68指向vq_model_pretrained_128_4096.pckl文件名即对应 128 维 × 4096 项的码本并带 SHA-256 校验和。on_train_start钩子中先检查预处理变换里不允许出现NormalizeDSR 输入不做 ImageNet 归一化否则直接抛ValueError然后调用prepare_pretrained_model()将权重下载到预训练权重目录下的dsr/子目录并由DsrModel.load_pretrained_discrete_model_weights以weights_onlyTrue载入。集成测试 tests/integration/data/test_download_urls.py 也会定期校验该下载 URL 的有效性。默认预处理器由configure_pre_processor给出仅Resize((256, 256), antialiasTrue)无归一化见 lightning_model.py#L336-L352。3.2 优化器与阶段切换由于三阶段共用一个 Lightning 训练循环Dsr采用手动优化self.automatic_optimization False并返回两个优化器optimizer_dlr2e-4Adam负责物体专属重建网络、两个 subspace 模块、异常检测模块并配一个StepLR调度器在第0.8 * second_phase步将学习率乘以 0.1optimizer_ulr2e-4Adam只负责上采样模块。其中second_phase int(num_steps * upsampling_train_ratio)num_steps取max_steps / len(train_dataloader)与max_epochs的较大者见 lightning_model.py#L148-L184。training_step根据current_epoch second_phase自动切换两条训练路径见 lightning_model.py#L234-L274第二阶段DsrAnomalyGenerator.augment_batch生成潜空间异常掩码 →DsrModel(input_image, anomaly_mask)走特征级缺陷分支 → 计算DsrSecondStageLoss→ 仅optimizer_d更新第三阶段PerlinAnomalyGenerator在图像级叠加 smudge 类污点 →DsrModel(input_image)走推理式前向 → 计算DsrThirdStageLoss→ 仅optimizer_u更新日志会打印 Now training upsampling module.。trainer_arguments固定返回{num_sanity_val_steps: 0}避免训练前额外跑验证。4. 特征级异常是如何“凭空”生成的这是 DSR 最有特色的部分。第二阶段的潜空间异常掩码由 src/anomalib/models/image/dsr/anomaly_generator.py 中的DsrAnomalyGenerator生成以p_anomalous0.5的概率决定该样本是否注入异常正常样本使用全零掩码异常掩码用 Perlin 噪声生成在[2^0, 2^6)内随机取对数尺度作为噪声频率阈值 0.5 二值化再叠加-90°~90°随机旋转最终得到(1, H, W)的二值掩码见 anomaly_generator.py#L52-L90。拿到掩码后DiscreteLatentModel.generate_fake_anomalies_joined见 torch_model.py#L1155-L1208在特征层面替换量化向量计算每个位置特征到码本全部向量的距离取“最不像正常”的前topk个码字topk 由异常强度决定且额外跳过最近的 5% 码字以避免选到自身随机挑一个作为该位置的“缺陷嵌入”再按掩码做mask * random_embeddings (1 - mask) * embeddings的混合。DiscreteLatentModel.forward还会以随机掩码use_both / use_lo / use_hi决定把缺陷注入 lo、hi 或两者见 torch_model.py#L1295-L1317从而覆盖三种污染情形。这就是文档所说“利用码本在量化表征上生成近分布缺陷”的源码落点。第三阶段则换成图像级模拟PerlinAnomalyGenerator位于anomalib.data.utils.generators.perlin在像素上叠加 smudge专门训练上采样模块使其对真实感污点更鲁棒。5. 损失函数第二阶段与第三阶段实现见 src/anomalib/models/image/dsr/loss.pyDsrSecondStageLoss四项之和——hi 特征 MSEl2(recon_feat_hi, embedding_bot)约束 subspace 模块把受污染 hi 特征重投影回正常量化表征lo 特征 MSEl2(recon_feat_lo, embedding_top)同理图像重建 MSEl2(input_image, obj_spec_image) * 10权重放大 10 倍分割 Focal Losskornia.losses.FocalLoss(alpha1)低分辨率预测掩码对真实下采样后的异常掩码。DsrThirdStageLoss仅一项 Focal Loss对上采样模块的全分辨率输出与图像级 smudge 真值掩码计算。Focal Loss 用于缓解分割任务中正常/异常像素的类别不平衡MSE 项则保证 subspace 重投影与重建网络的“去缺陷”能力。6. 实际使用配置与训练命令6.1 模型配置文件仓库提供的官方配置 examples/configs/model/dsr.yaml 内容如下model: class_path: anomalib.models.Dsr init_args: latent_anomaly_strength: 0.2 upsampling_train_ratio: 0.7 # PL Trainer Args. Dont add extra parameter here. trainer: max_epochs: 700即按max_epochs: 700训练前 490 个 epoch0.7 * 700训练重建/subspace/异常检测模块后 210 个 epoch 训练上采样模块。学习率退火步约为第 392 epoch0.8 * 490后降为原来的 0.1。6.2 训练与推理模型目录下的 src/anomalib/models/image/dsr/README.md 给出的标准训练命令为anomalib train --model Dsr --data MVTecAD --data.category category也可以在 Python API 中直接实例化from anomalib.models.image import Dsr model Dsr( latent_anomaly_strength0.2, upsampling_train_ratio0.7, ) print(model.trainer_arguments) # {num_sanity_val_steps: 0}需要注意的适用前提训练开始时会自动下载约一次性的 ImageNet 预训练 VQ 权重见第 3.1 节首次运行需联网输入数据经Resize(256, 256)预处理且不做归一化自定义 transform 中若混入Normalize会直接报错。验证阶段validation_step直接调用self.model(batch.image)获得pred_score与anomaly_map再经batch.update合并回批次。7. 小结DSR 在 Anomalib 中定位为Segmentation 型、one-class 学习的量化特征算法核心卖点是用 ImageNet 预训练码本在特征空间合成“近分布”缺陷从而摆脱对外部缺陷数据的依赖源码上由DsrModel五大网络模块 冻结的DiscreteLatentModel、Dsr三阶段训练调度与预训练权重管理、DsrAnomalyGenerator/PerlinAnomalyGenerator两阶段异常模拟与两个 Focal/MSE 组合损失构成完整闭环上手时只需记住两个超参数latent_anomaly_strength潜空间缺陷强度下界默认 0.2与upsampling_train_ratio第二/三阶段切换比例默认 0.7再配合官方 YAML 中的max_epochs: 700即可完成训练。想进一步了解量化细节可直接阅读 torch_model.py 中VectorQuantizer.forward与generate_fake_anomalies_joined两处实现模型注册入口见 src/anomalib/models/image/init.py。【免费下载链接】anomalibAn anomaly detection library comprising state-of-the-art algorithms and features such as experiment management, hyper-parameter optimization, and edge inference.项目地址: https://gitcode.com/GitHub_Trending/an/anomalib创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表