ProbEdge-LiteSeg 是一个面向医学图像辅助诊断的轻量化病灶分割与量化部署评估框架。针对轻量化模型在病灶边缘细节刻画和 INT8 量化部署中存在的双重痛点，本框架提出基于概率空间形态边缘约束的分割方案，并构建了从模型训练、结构重参数化导出到 ONNX Runtime 量化部署的完整评估闭环。
核心目标：在严格的资源约束（4核CPU/4GB内存，384x384输入下，端到端延迟 < 3000ms）下，输出高保真的病灶掩膜、面积占比与重心方位特征。此处端到端延迟包含：模型前向推理（含 Sigmoid）、阈值二值化与开运算去噪，以及连通域几何特征提取。
核心发现：实验证实，基于重参数化的现代轻量架构（如 FastViT）在 ONNX Runtime CPU 后端下存在严重的 INT8 算子降级回退风险，导致精度崩塌与加速比倒挂。本框架通过严谨的量化可行性闭环校验机制，揭示了这一工程隐患，并确立了 FastViT FP32 极速合规 + MobileNetV2 INT8 兜底的分级部署策略。
-
概率空间形态边缘约束：在连续概率场上提取 Sobel 梯度并进行
$\sqrt{32}$ 归一化，结合形态学预计算边缘标签与带权 BCE，显式优化边界锐利度，规避数值溢出风险。 - 严谨的样本级评估体系：纯 PyTorch/Numpy 向量化实现样本级 Dice/IoU 计算，严格区分真阴性（TN）与假阴性（FN），引入正样本 Dice (Pos-Dice)、面积相对误差与重心偏移距离，甄别退化解。
-
架构感知的部署管线：彻底摒弃 Apple 官方残缺源码，通过
timm统一管理预训练权重与重参数化逻辑；采用copy.deepcopy权重隔离导出机制，并封装 Sigmoid 至计算图以支持感知概率分布的 INT8 校准。 - 量化可行性闭环校验：内置 INT8 加速比与精度衰减的硬性阈值阻断机制，自动决策量化的工程可行性，防范“量化即崩塌”的线上事故。
基于 ISIC 2018 Task 1 数据集的严格控制变量实验与云端（4核CPU/4GB内存）实测，得出以下结论：
| 实验编号 | 编码器架构 | 边缘约束 | PyTorch Pos-Dice | FP32 P95延迟 | INT8 P95延迟 | INT8 加速比 | 量化校验结果 | 部署决断 |
|---|---|---|---|---|---|---|---|---|
| Exp-1 | ResNet18 | 无 | 0.8925 | 144.08 ms | 71.74 ms | 2.01x | 量化稳健 | 标准卷积基线 |
| Exp-2 | FastViT-T8 | 无 | 0.8970 | 71.73 ms | 126.53 ms | 0.57x | 精度崩塌/加速倒挂 | 最终选型 (FP32) |
| Exp-3 | FastViT-T8 | 有 | 0.8878 | 83.21 ms | 133.43 ms | 0.62x | 精度崩塌/加速倒挂 | 边缘约束反噬 |
| Exp-4 | MobileNetV2 | 有 | 0.8761 | 119.93 ms | 58.95 ms | 2.03x | 量化稳健 | 兜底容灾 (INT8) |
关键决策分析：
- 架构逆袭：Exp-2 (FastViT) 以仅 ~4.5M 参数量，在 PyTorch 精度上反超 ~11M 的 ResNet18，且 FP32 推理延迟极低（71.73ms），远低于 3000ms 红线。
- INT8 量化灾难：FastViT 的重参数化结构在 ONNX Runtime 下遭遇算子降级回退，INT8 量化后 Dice 跌破 0.2，且速度不升反降（加速比 0.57x）。此现象证明了量化可行性闭环校验的绝对必要性。
- 边缘约束反噬：在当前超参下，边缘损失的多目标优化对轻量架构的区域收敛造成了负优化冲突（Exp-3 vs Exp-2）。
- 最终选型：选用 Exp-2 (FastViT 无Edge) 的 FP32 模型部署，兼顾极速与高精度；若未来输入分辨率暴增导致 FP32 延迟超标，触发阻断机制，回退至 Exp-4 (MobileNetV2 INT8) 兜底。
本项目严格分离训练环境与部署环境，以适应本地高性能计算与受限于低资源推理的物理隔离需求。
pip install -r requirements/train_requirements.txt(注：需根据本地 CUDA 版本配置 PyTorch，onnxruntime 为本地导出校验与量化校准必需)
pip install -r requirements/deploy_requirements.txt(注：纯 ONNX Runtime 推理，无需 PyTorch，规避 OOM 风险)
ProbEdge-LiteSeg/
├── configs/ # 正交实验配置文件 (YAML)
├── datasets/ # 数据集加载与形态学边缘标签在线生成逻辑
├── losses/ # 概率空间形态边缘约束损失及联合损失定义
├── models/ # 自定义编码器注入 (基于 timm 的 FastViT) 及动态解码器构建
├── utils/ # 指标计算、种子固定等基础工具类
├── scripts/ # 训练、导出、量化、测速执行脚本
│ ├── train_research.py # 统一训练入口 (双轨早停)
│ ├── export_onnx.py # 深拷贝隔离与重参数化 ONNX 导出
│ ├── quantize_local.py # 本地 INT8 静态量化 (防范云端 OOM)
│ ├── eval_onnx.py # 本地 ONNX 精度对齐校验
│ └── benchmark_onnx.py # 云端端到端延迟剖析与量化闭环校验
├── requirements/ # 环境依赖文件
├── onnx_weights/ # 导出产物存放目录 (不入库)
├── pytorch_weights/ # 训练权重存放目录 (不入库)
└── logs/ # TensorBoard 日志目录 (不入库)
下载 ISIC 2018 Task 1 训练集图像与掩膜，解压至项目根目录。数据加载模块已内置标准化与 3x3 椭圆核形态学边缘标签在线生成。
采用控制变量法设计，四组对照实验需独立执行：
python scripts/train_research.py --config configs/exp1_resnet18.yaml
python scripts/train_research.py --config configs/exp2_fastvit_noedge.yaml
python scripts/train_research.py --config configs/exp3_fastvit_edge.yaml
python scripts/train_research.py --config configs/exp4_mobilenetv2_edge.yaml深拷贝隔离权重，融合多分支结构，并封装 Sigmoid 至计算图（Opset 15）：
python scripts/export_onnx.py --config configs/exp2_fastvit_noedge.yaml针对受限云端内存不足 (OOM) 导致量化崩溃的问题，提供本地量化脚本，采用 Percentile 校准策略保真 Sigmoid 输出：
python scripts/quantize_local.py将生成的 FP32 与 INT8 模型上传至目标云端环境，执行端到端延迟剖析与量化可行性闭环校验：
python3 scripts/benchmark_onnx.py \
--onnx_dir ./ \
--val_img_dir <验证集图像路径> \
--mask_dir <验证集掩膜路径> \
--img_size 384本项目采用正交分组与严格的级联决策准则，剥离架构演进与边缘约束的独立贡献。
| 实验编号 | 编码器架构 | 解码器策略 | 损失函数配置 | 实验假设与验证目的 |
|---|---|---|---|---|
| Exp-1 | ResNet18 | 轻量化解码器 | BCE + Dice | 标准卷积量化稳定性基线 |
| Exp-2 | FastViT-T8 | 轻量化解码器 | BCE + Dice | 重参数化现代架构精度与速度收益验证 |
| Exp-3 | FastViT-T8 | 轻量化解码器 | BCE + Dice + Edge | 概率空间形态边缘约束对边界锐利度与几何特征的独立增益验证 |
| Exp-4 | MobileNetV2 | 轻量化解码器 | BCE + Dice + Edge | NAS 架构在 INT8 量化下的算子稳健性容灾兜底验证 |
级联决策原则：INT8 P95延迟 < 3000ms → 加速比 > 1.2x → Pos-Dice 下降 < 2% → 几何特征退化低于阈值 → 精度优先 → 边缘可视化决断。
Q1: 为什么 ONNX 导出时必须使用深拷贝 (copy.deepcopy) 并封装 Sigmoid？
A: FastViT 等重参数化架构在导出前必须调用 reparameterize() 将多分支融合为单路卷积，此操作对权重具有不可逆的破坏性。深拷贝可保护训练期模型权重不被污染；将 Sigmoid 封装入计算图，有利于 INT8 校准器对概率输出分布进行整体感知。
Q2: 为什么使用 timm 库加载 FastViT 而不是 Apple 官方源码？
A: Apple 官方开源的 FastViT 代码库阉割了预训练权重的自动下载接口，手动加载极易出错。timm 作为权威模型库，完美收录了 FastViT 并提供一键预训练权重下载与 features_only 接口，工程稳健性远高于本地引入残缺源码。
Q3: 为什么 FastViT (Exp-2/3) 的 INT8 量化会失败（加速比倒挂、精度崩塌）？ A: 这是重参数化架构在 ONNX Runtime CPU 后端的典型工程隐患。融合后的特定卷积算子无法被 ORT 的 INT8 执行引擎正确识别与加速，触发了算子降级回退至 FP32 或低效实现，不仅导致精度断崖式下跌，还引入了额外的类型转换开销致使速度变慢。这正是本框架内置量化可行性闭环校验的价值所在。
Q4: 云端受限环境（4GB内存）跑静态量化校准时 OOM 怎么办？
A: ONNX 的静态量化在校准阶段需缓存大量中间激活值，极易 OOM。本框架的应急预案是：将量化步骤移至内存充足的本地环境执行（scripts/quantize_local.py），生成 INT8 模型后，再上传至云端仅做推理测速与精度校验。
本项目基于 MIT License 开源。
- 数据集提供：ISIC 2018: Task 1 - Lesion Boundary Segmentation
- 基础框架与模型：segmentation-models-pytorch, Apple FastViT, PyTorch Image Models (timm)