ViT_Fast / EXPERIMENTS.md

1999xia

Upload folder using huggingface_hub

54ee1eb verified about 2 months ago

preview code

Raw

History Blame Contribute Delete

15.3 kB

Miss Patch — Pre-Tokenization Patch Selection for Vision Transformers

在 Vision Transformer（ViT）前插入一个 Router/Filter 模块，动态过滤不重要的 image patches，减少计算量，同时尽量保持分类精度。

项目文件

文件	说明
`models.py`	所有模型定义，见下方"模型类"表格
`train.py`	通用训练脚本（支持 baseline、patch selection、MAE finetune 等）
`train_patch_selection_mae.py`	MAE patch selection 专用训练脚本（含蒸馏 router 加载）
`train_router_distill.py`	Attention Distillation 训练脚本
`test_patch_selection_b16.py`	旧 Gumbel patch selection 测试脚本
`test_blur_downsample.py`	图片模糊/降采样测试脚本
`test_stride_patches.py`	不同 stride 下 patch 数量 vs 精度测试脚本
`datasets.py`	数据集加载器：CIFAR-10/100, Oxford Pets, Food-101, Tiny-ImageNet, DTD, Flowers-102, Stanford Cars
`checkpoints/`	所有模型权重
`logs/`	训练日志
`SHARED_MEMORY.md`	跨 Claude Code 会话共享记忆

models.py — 模型类

类名	说明
`GumbelSelection`	Gumbel-Softmax 可微分选择模块（含 STE + 退火温度）
`SemanticRouter`	语义 Router（MLP: D→D/2→1 + GELU + LayerNorm）
`PatchSelectionViT`	Patch Selection ViT（Router + 可微分 Top-K + Gumbel 噪声）
`RandomPruneViT`	随机丢弃 patch 的 baseline
`MAEDecoder`	MAE 解码器（4 blocks, 512-dim transformer）
`MAEViT`	标准 MAE（mask 75% patches + 重建）
`MAEPatchSelectionViT`	MAE + Patch Selection（Router + Differentiable Top-K + 轻量 encoder + 10 blocks backbone + MAE Decoder）

models.py — create_model() 可用 model_name

model_name	对应模型	说明
`swin_tiny`	Swin-Tiny	Swin Transformer baseline
`patch_selection_vit`	PatchSelectionViT	Gumbel 方案（旧）
`patch_selection_vit_b16`	PatchSelectionViT	ViT-B/16 IN-21K + Gumbel
`patch_selection_vit_b16_in1k`	PatchSelectionViT	ViT-B/16 IN-1K + Gumbel
`random_prune_vit`	RandomPruneViT	随机丢弃 50% baseline
`mae_vit`	MAEViT	标准 MAE 预训练
`mae_patch_selection_vit_b16`	MAEPatchSelectionViT	MAE + 蒸馏 Router（推荐）

使用方法

1. 训练 Baseline（全量 ViT-B/16）

# CIFAR-100 (IN-21K pretrained)
python train.py --model vit_b16 --dataset cifar100 --lr 3e-5 --epochs 100 --batch_size 128

# Oxford Pets
python train.py --model vit_b16 --dataset oxford_pets --lr 3e-5 --epochs 100 --batch_size 32

# Food-101
python train.py --model vit_b16 --dataset food101 --lr 3e-5 --epochs 30 --batch_size 32

2. 训练 MAE Patch Selection（+ 蒸馏 Router）

两步走：

Step 1: 训练蒸馏 Router（从教师 ViT 学习注意力分布）

python train_router_distill.py --dataset cifar100 --gpu 4
# 可选参数：
#   --batch_size 64     # 批量大小（默认 64）
#   --lr 1e-4           # 学习率（默认 1e-4）
#   --epochs 5          # 训练轮数（默认 5）
#   --weight_decay 0.05 # 权重衰减

# 输出：checkpoints/router_distill_{dataset}/router.pth

Step 2: 用蒸馏 Router 初始化，训练完整模型

python train_patch_selection_mae.py --dataset cifar100 --gpu 4 \
  --router_path ./checkpoints/router_distill_cifar100/router.pth

# 可选参数：
#   --keep_ratio 0.5    # 保留 patch 比例（默认 0.5，推荐 0.75）
#   --batch_size 32     # 批量大小（默认 32）
#   --accum 4           # 梯度累积步数（默认 4，有效 batch = 128）
#   --lr 3e-5           # 学习率（默认 3e-5）
#   --weight_decay 0.05 # 权重衰减
#   --label_smoothing 0.1
#   --mse_start 1.0     # MSE 损失起始权重
#   --mse_end 0.1       # MSE 损失结束权重（cosine anneal）
#   --decoder_dim 512   # MAE decoder 维度
#   --decoder_depth 4   # MAE decoder 层数
#   --router_path PATH  # 蒸馏 router 权重路径（可选）

不加 --router_path 就是随机初始化 Router（原始 MAE 方案）。

3. 训练 Gumbel Patch Selection（旧方案，不推荐）

python train.py --model patch_selection_vit_b16 --dataset cifar100

4. 测试图片预处理（模糊/降采样）

python test_blur_downsample.py --dataset cifar100 --gpu 5
# 测试 Gaussian blur、降采样等预处理对精度的影响（不减少 token 数量）

5. 测试 Stride-based Patch 减少

python test_stride_patches.py --dataset cifar100 --gpu 5
# 测试不同 stride 下 patch 数量 vs 精度（直接减少 token 数量）

6. 降采样图片训练（减小输入分辨率）

# 112×112 → 49 patches (75% reduction)
python test_downsample_train.py --dataset cifar100 --image_size 112 --gpu 3

# 168×168 → 100 patches (49% reduction)
python test_downsample_train.py --dataset cifar100 --image_size 168 --gpu 5

# 可选参数：
#   --batch_size 128   # 批量大小（默认 128）
#   --lr 3e-5          # 学习率（默认 3e-5）
#   --epochs 100       # 训练轮数（默认使用数据集默认值）
#   --label_smoothing 0.1
#   --weight_decay 0.05
#   --num_workers 4

# 输出：checkpoints/{dataset}_vit_b16_img{size}/best_model.pth

7. 通用训练脚本 train.py 全部参数

python train.py \
  --model vit_b16                # 模型名
  --dataset cifar100             # 数据集
  --batch_size 128               # 批量大小
  --epochs 100                   # 训练轮数
  --lr 3e-5                      # 学习率
  --weight_decay 0.05            # 权重衰减
  --keep_ratio 0.5               # patch 保留比例
  --selection_mode topk          # 选择模式
  --adaptive_alpha 0.5           # 自适应 alpha
  --patch_size 16                # patch 大小
  --patch_stride 16              # patch stride
  --num_workers 10               # 数据加载线程
  --use_randaugment              # 使用 RandAugment
  --label_smoothing 0.0          # 标签平滑
  --device auto                  # 设备
  --save_dir ./checkpoints       # 保存目录
  --pretrained                   # 使用预训练权重
  --load_mae PATH                # 加载 MAE 预训练权重
  --log_dir ./logs               # 日志目录
  --drop_path 0.0                # DropPath 率
  --mixup 0.0                    # MixUp alpha
  --cutmix 0.0                   # CutMix alpha
  --ema_decay 0.0                # EMA 衰减率
  --clip_grad 1.0                # 梯度裁剪
  --warmup_epochs 0              # 预热轮数

支持的数据集

datasets.py 提供：

get_cifar10_loader — CIFAR-10 (10 classes, 50K train)
get_cifar100_loader — CIFAR-100 (100 classes, 50K train)
get_oxford_pets_loader — Oxford Pets (37 classes, ~3.6K train)
get_food101_loader — Food-101 (101 classes, 68K train)
get_tiny_imagenet_loader — Tiny-ImageNet (200 classes, 100K train)
get_dtd_loader — DTD (47 classes, texture)
get_flowers102_loader — Flowers-102 (102 classes)
get_stanford_cars_loader — Stanford Cars (196 classes)

项目架构

Baseline 架构（全量 ViT-B/16 IN-21K）

Image (224×224) → Patch Embed (196 patches, 768-dim) → 12 ViT-B blocks → CLS → Classifier

MAE Patch Selection 架构（当前最佳）

Image → Patch Embed + Pos Embed → Router (MLP 768→384→1) → Sigmoid STE Top-K (keep 75%)
  → 选中 patches (147个) + CLS → 轻量 Encoder (前 2 个 ViT-B block)
    → 主干 (后 10 个 ViT-B block) → CLS → CE Loss
    → MAE Decoder (4 blocks, 512-dim) → 重建丢弃 patches → MSE Loss (λ: 1.0→0.1)

Attention Distillation 流程

教师 ViT-B/16（冻结）→ 提取最后 block 的 CLS→patch 注意力 → 注意力分数 (B, 196)
MLP Router（学生）→ 从 patch embedding 预测重要性分数 → MSE Loss
→ 训练好的 Router 权重→ 初始化 MAEPatchSelectionViT → 端到端微调

实验结果汇总

主数据集：Oxford-IIIT Pets（37 类，高分辨率 ~200-1000px）

全量对比（相同条件：batch=16, accum=2, lr=3e-5, 100 epochs）：

方法	保留	Test Acc	vs Baseline
Baseline（全量 ViT-B/16）	100%	91.91%	—
降采样 168x168	51% patches	90.65%	-1.26%
灰度图（Grayscale）	亮度 only	90.68%	-1.23%
MAE + 从头训练 Router	50% patches	88.80%	-3.11%
降采样 112x112	25% patches	86.64%	-5.27%
MAE + 预热 Router	50% patches	85.85%	-6.06%

结论：

降采样 168x168 和灰度图效果最好（掉 ~1.2%），且不需要额外模块
MAE + 从头训练 Router（88.80%）比预热过的 Router（85.85%）更好——在小数据集上，教师蒸馏信号太弱（重叠率仅 51.16%≈随机），预热反而有害
降采样比所有学习型方案更简单、效果更好

验证数据集

Food-101（101 类，细粒度菜肴，~300-512px）

方法	Test Acc	vs Baseline
Baseline（batch=128）	91.37%	—
降采样 168x168	89.87%	-1.50%
降采样 112x112	85.96%	-5.41%
MAE + 预热 Router（50%，CIFAR）	89.52%	-1.85%（参考）

DTD（47 类，纹理分类，~300-400px）

方法	Test Acc	vs Baseline
Baseline	80.85%	—
降采样 168x168	74.20%	-6.65%

跨数据集对比：降采样 168x168

数据集	特点	Baseline	168x168	下降幅度
CIFAR-100	粗粒度，原生 32x32	91.69%	91.56%	-0.13%（假象）
Food-101	细粒度菜肴	91.37%	89.87%	-1.50%
Oxford Pets	猫狗品种，花纹纹理	~93.3%	90.65%	-3.16%
DTD	纯纹理分类	80.85%	74.20%	-6.65%

规律：分类粒度越细、越依赖高频纹理，降采样伤害越大。

旧实验（CIFAR-100，仅供参考）

CIFAR-100 原生仅 32x32，降采样结果不可靠，但 Router 类实验（不改变图片大小）有效。

方法	保留	CIFAR-100	vs BL
Baseline	100%	91.69%	—
Gumbel Selection	50%	87.18%	-4.51%
MAE + 随机 Router	50%	88.25%	-3.44%
MAE + 蒸馏 Router	50%	89.07%	-2.62%
MAE + 蒸馏 Router	75%	91.10%	-0.59%

Gumbel 方案最差 — Gumbel 噪声导致 Router 梯度几乎消失，无法有效学习 patch 重要性
MAE 重建损失有帮助 — 让 Router 通过"能否重建丢弃 patch"来学习，比纯分类信号好（+1%）
Attention Distillation 最佳 — 用教师 ViT 的 CLS 注意力预训练 Router，再端到端微调
- CIFAR-100: 88.25% → **89.07%**（+0.82%）
- Food-101: ~87.97% → **89.52%**（+1.55%）
数据集大小重要 — Oxford Pets 仅 ~3.6K 训练图像，蒸馏信号太弱，反而下降
CLS token flow 关键 bug — 早期实现中 CLS 只走 10/12 block（前 2 block 浪费），修复后 epoch 1 从 80.63%→84.69%
75% keep ratio 是更优平衡点 — 只掉 0.59%（vs 50% 掉 2.62%），推荐使用
降低 token 质量 vs 减少 token 数量 — 模糊/降采样图片几乎不影响精度（-0.14~-0.57%），但计算量完全不变；减少 token 数量才真正降低计算量
Naive stride 减少 token 效果差且实用性低 — stride=32 减少 75% token 时掉 21%，远不如学习型 Router 或降采样训练

检查点文件

checkpoints/
├── cifar100_mae_patchsel_b16_keep50/
│   └── best_model.pth          —— 89.07% (MAE + 蒸馏 Router, keep 50%)
├── cifar100_mae_patchsel_b16_keep75/
│   └── best_model.pth          —— 91.10% (MAE + 蒸馏 Router, keep 75%) ★推荐
├── food101_mae_patchsel_b16_keep50/
│   └── best_model.pth          —— 89.52% (MAE + 蒸馏 Router, keep 50%)
├── oxford_pets_mae_patchsel_b16_keep50/
│   └── best_model.pth          —— 85.99% (MAE + 蒸馏 Router, keep 50%)
├── cifar100_patchsel_b16_keep50/
│   └── best_model.pth          —— 87.18% (Gumbel, CIFAR-100)
├── food101_patchsel_b16_keep50/
│   └── best_model.pth          —— 85.77% (Gumbel, Food-101)
├── oxford_pets_patchsel_b16_keep50/
│   └── best_model.pth          —— 87.00% (Gumbel, Oxford Pets)
├── cifar100_random_prune_vit_keep50_topk/
│   └── best_model.pth          —— 48.13% (随机丢弃 50%)
├── cifar100_vit_b16_ft/
│   └── best_model.pth          —— 88.75% (全量微调, CIFAR-100)
├── cifar100_vit_small/
│   ├── best_model.pth          —— 84.83% (ViT-S/16 baseline)
│   └── checkpoint_epoch*.pth   —— 每 10 epoch 检查点
├── cifar100_patch_selection_vit_keep50_topk/
│   ├── best_model.pth          —— 73.01% (旧 Gumbel 实验)
│   └── checkpoint_epoch*.pth   —— 每 10 epoch 检查点
├── vit_b16_cifar10_best.pth    —— CIFAR-10 baseline
├── vit_b16_cifar10_ft/
│   └── best_model.pth          —— CIFAR-10 微调
├── dtd_vit_b16_in21k/
│   └── best_model.pth          —— DTD 80.85%
├── flowers102_vit_b16_in21k/
│   └── best_model.pth          —— Flowers-102 100.00%
├── food101_vit_b16_in21k/
│   └── best_model.pth          —— Food-101 86.68% (val)
├── oxford_pets_vit_b16_ft/
│   └── best_model.pth          —— Oxford Pets 全量微调
├── oxford_pets_vit_b16_in21k/
│   └── best_model.pth          —— Oxford Pets IN-21K 95.65% (val)
├── mae_cifar100_mask75/
│   ├── mae_best.pth            —— MAE 预训练最佳
│   ├── mae_encoder_final.pth   —— MAE encoder 最终
│   ├── mae_epoch50.pth         —— 50 epoch 检查点
│   └── mae_epoch100.pth        —— 100 epoch 检查点
├── router_distill_cifar100/
│   └── router.pth              —— 蒸馏 Router (CIFAR-100, 55.76% overlap)
├── router_distill_food101/
│   └── router.pth              —— 蒸馏 Router (Food-101, 59.95% overlap)
└── router_distill_oxford_pets/
    └── router.pth              —— 蒸馏 Router (Oxford Pets, 51.16% overlap)

超参数（推荐）

所有 MAE patch selection 实验统一使用：

参数	值
Backbone	ViT-B/16 IN-21K pretrained (`vit_base_patch16_224.augreg_in21k`)
Batch size	32
Gradient accumulation	4（effective batch = 128）
Learning rate	3e-5
Weight decay	0.05
Label smoothing	0.1
Scheduler	CosineAnnealingLR
Gradient clip	1.0
Keep ratio	0.75（147/196 patches，推荐）或 0.5（98/196 patches）
MSE weight	1.0 → 0.1（cosine anneal）
Decoder dim	512
Decoder depth	4
Epochs	CIFAR-100: 100, Oxford Pets: 100, Food-101: 30

最后更新: 2026-05-04