SG-JEDNet-Tiny

轻量化水下显著性目标检测模型 | 参数量 2.49M | USOD10K S-measure 0.7575

面向水下场景的 RGB-D 显著性目标检测，以最小参数量和计算代价完成精确检测。

模型架构

RGB (3×256×256) + Depth (1×256×256)
        │
   MobileNetV3-Small Backbone
   ├── stage1: (24, 32, 32)  ← 1/8  → SG-IEB merge1
   ├── stage2: (40, 16, 16)  ← 1/16 → SG-IEB bottleneck
   └── stage3: (576, 8, 8)   ← 1/32 → SDB init pass

SDB-Tiny (检测分支):
  init:          stage3 → proj(576→128) → decode1
  enhanced:      SDBEncoder(enhanced_img) → decode1
  decode1-4:     1/32 → 1/16 → 1/8 → 1/4 → 1/1
  预测头:       sal_1_16, sal_1_8, sal_1_4, sal_full

模块说明

模块	功能	参数
Backbone	MobileNetV3-Small，提取 RGB/Depth 多尺度特征	0.93M
SDB-Tiny	检测分支，5层深度可分离卷积解码器	1.18M
SG-IEB	显著性引导图像增强分支	0.02M
IMFE-Lite	轻量多模态融合，64维深度编码	0.20M

详细架构见 ARCHITECTURE.md。

性能

指标	v4 (2.49M)	v1 (3.96M)
S-measure (Sα)	0.7575	0.7768
F-measure (Fβ^ω)	0.7741	0.8001
MAE	0.0465	—
Epochs to Best	27	47

环境配置

conda create -n usod python=3.10
conda activate usod
pip install torch>=2.0.0 torchvision>=0.15.0 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

硬件要求：NVIDIA GPU 8GB+（RTX 3070 测试通过），输入 256×256，batch_size=8。

数据集

使用 USOD10K 数据集，目录结构：

data/USOD10K/
├── images/{train,val,test}/
├── masks/{train,val,test}/
├── depths/{train,val,test}/
└── edges/{train,val,test}/

训练

# 快速验证（短配置）
python scripts/train.py --config configs/short_train.yaml

# 正式训练 v4 纯检测模式
python scripts/train.py --config configs/simplified_detect.yaml --auto-resume

测试

python scripts/test.py \
    --config configs/simplified_detect.yaml \
    --checkpoint ./outputs_v4_simplified/checkpoints/best_model.pth \
    --data_split test --save_predictions

项目结构

sgjednet-tiny/
├── configs/          # 训练配置
├── models/           # 模型定义
├── engine/           # 训练引擎
├── data/             # 数据加载
├── utils/            # 工具函数
├── scripts/          # 入口脚本
├── ARCHITECTURE.md   # 架构详述
├── VERSION_LOG.md    # 版本训练记录
└── RESUME_TRAINING.md # 断点续训指南

关键设计

1. IMFE 64ch 独立控制 — 深度图语义简单，64维编码足够
2. Decoder 全 128ch 无瓶颈 — 保留全分辨率细节
3. 语义反馈 @1/8 (32×32) — 相比 1/32 精度提升 16 倍
4. Backbone BN 冻结 — m.eval(), weight.requires_grad = False
5. 纯检测模式 — 增强分支梯度回传干扰 backbone，关闭后效果更好

训练经验

• v4 起步快但撞墙早：epoch 0 S=0.6357，epoch 27 即触顶
• 深度图至关重要：去掉深度图 S 从 0.75 降至 0.53
• ImageNet 预训练比随机初始化高 0.09 S-measure
• 端到端训练优于两阶段（0.757 vs 0.716）
• Loss 持续降但 S 不涨 → 过拟合，考虑更早早停

License

MIT