# vit_jax

**Repository Path**: fmscole/vit_jax

## Basic Information

- **Project Name**: vit_jax
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-06-05
- **Last Updated**: 2026-06-25

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Vision Transformer (ViT) & Swin Transformer in JAX — CIFAR-10

纯 JAX 实现的 Vision Transformer 和 Swin Transformer，从零构建，无需 Flax/Haiku 等高层框架。涵盖监督训练、MAE 自监督预训练与微调的完整流程，并包含自定义 Flash Attention 内核。

> **新增**：`src/vit/swin_cifar10_cnn_stem.py` — 在 Swin Transformer 的 Patch Embedding（StemLayer）之前加入可学习的 **CNN PreStem** 局部特征预处理模块，专为 CIFAR-10 等小尺寸图像设计，参数量仅增加约 0.15M。

## 项目结构

```
vit_jax/
├── src/
│   ├── vit/                                    # 模型与训练脚本
│   │   ├── vit_cifar10_stax.py                 # ViT 监督训练
│   │   ├── vit_cifar10_mae_pretrain_stax.py    # MAE 预训练 (Stage 1)
│   │   ├── vit_cifar10_mae_finetune_stax.py    # MAE 微调 (Stage 2)
│   │   ├── swin_cifar10_stax.py                # Swin Transformer 监督训练
│   │   ├── swin_cifar10_cnn_stem.py            # Swin + CNN PreStem 监督训练
│   │   └── swin_encoder.py                     # Swin Transformer 编码器 / CNNPreStem
│   ├── transformers/                           # Transformer 基础模块
│   │   ├── transformer_flash_cudnn.py          # 基于 FlashAttention v1 的 Transformer
│   │   ├── transformer_flash_v2.py             # 基于 FlashAttention-2 的 Transformer
│   │   └── mae_utils.py                       # MAE 掩码/还原工具
│   ├── flash_attention/                        # 自定义 Flash Attention 内核
│   │   ├── flash_attention_v1.py               # FlashAttention (Q-outer loop)
│   │   └── flash_attention_v2.py               # FlashAttention-2 (KV-outer loop)
│   └── lib/                                    # 基础库
│       └── stax_plus.py                        # 扩展版 Stax 神经网络库
└── data/                                       # CIFAR-10 数据集（自动下载）
```

## 模型架构

### ViT

```
输入图像 (32×32×3)
    │
    ▼
Patch Embedding ─── Conv(embed_dim, 4×4, stride=4) → (batch, 64, embed_dim)
    │
    ▼
CLS Token ─── 可学习分类标记 → (batch, 65, embed_dim)
    │
    ▼
Positional Encoding ─── 可学习位置编码 → (batch, 65, embed_dim)
    │
    ▼
Transformer Encoder × N ─── Pre-LN + Flash MSA + FFN (GELU)
    │
    ▼
Classification Head ─── LayerNorm → [CLS] → Dense(num_classes)
```

| 参数 | ViT 默认值 | 说明 |
|------|-----------|------|
| `patch_size` | 4 | 32×32 → 64 个 patch |
| `embed_dim` | 384 | 模型维度 |
| `num_layers` | 7 | Encoder 层数 |
| `num_heads` | 6 | 注意力头数 |
| `mlp_dim` | 768 | FFN 中间层 |
| `dropout` | 0.2 | 训练时 Dropout |

### Swin Transformer

基于 Swin Transformer v1，包含 StemLayer（双层 Conv 下采样）、Window/Shifted-Window MSA、Relative Position Bias、PatchMerging 和 ConvEnhance。

```
输入 (B, 3, 32, 32)
  → StemLayer → (B, 96, 8, 8)
  → Stage 0 (dim=96, depth=2, heads=3, window=4)
  → PatchMerging → (B, 192, 4, 4)
  → Stage 1 (dim=192, depth=2, heads=6, window=4)
  → GlobalAvgPool → (B, 192)
  → Dense(768) → GELU → Dense(10)
```

| 参数 | Swin 默认值 |
|------|------------|
| `embed_dim` | 96 |
| `depths` | (2, 2) |
| `num_heads` | (3, 6) |
| `window_size` | 4 |
| `mlp_ratio` | 4.0 |
| 参数量 | ~1.38M |

### Swin + CNN PreStem

在原始 Swin StemLayer 之前插入多层 `stride=1` 的 Conv2d + GELU，作为可学习的局部特征预处理模块（`CNNPreStem`），增强小图像的局部特征提取。

```
输入 (B, 3, 32, 32)
  → CNNPreStem [(32,3), (64,3), (128,3)] → (B, 128, 32, 32)
  → StemLayer → (B, 96, 8, 8)
  → Stage 0 (dim=96, depth=2, heads=3, window=4)
  → PatchMerging → (B, 192, 4, 4)
  → Stage 1 (dim=192, depth=2, heads=6, window=4)
  → GlobalAvgPool → (B, 192)
  → Dropout → Dense(768) → GELU → Dropout → Dense(10)
```

| 参数 | Swin + CNN PreStem 默认值 |
|------|------------|
| `cnn_layers` | `[(32, 3), (64, 3), (128, 3)]` |
| `embed_dim` | 96 |
| `depths` | (2, 2) |
| `num_heads` | (3, 6) |
| `window_size` | 4 |
| `mlp_ratio` | 4.0 |
| `dropout_rate` | 0.1 |
| 参数量 | ~1.53M |

## 训练范式

### 1. ViT 监督训练

从随机初始化训练 ViT 分类器。

```bash
python src/vit/vit_cifar10_stax.py
```

**训练配置：**
- 500 epochs, batch_size=128
- AdamW (lr=3e-4, weight_decay=0.1)
- Warmup 10 epochs + Cosine decay
- MixUp (α=0.8) + RandAugment + Label Smoothing (0.1)
- Gradient Clipping (max_norm=1.0)

### 2. Swin 监督训练

从随机初始化训练 Swin Transformer 分类器（适配 32×32 小图）。

```bash
python src/vit/swin_cifar10_stax.py
```

**训练配置：** 与 ViT 监督训练相同。

### 3. Swin + CNN PreStem 监督训练

从随机初始化训练带 CNN 预处理的 Swin Transformer 分类器，适用于 CIFAR-10 等小尺寸图像。

**设计动机**：CIFAR-10 仅 32×32，Swin 的 StemLayer 直接下采样到 8×8 会损失部分细粒度局部信息。`CNNPreStem` 在 StemLayer 前用若干 `stride=1` 的 Conv2d + GELU 提取局部特征，分辨率保持 32×32，几乎不增加计算开销，同时提升最终分类精度。

```bash
python src/vit/swin_cifar10_cnn_stem.py
```

**训练配置：**
- 500 epochs, batch_size=128
- AdamW (lr=3e-4, weight_decay=0.1)
- Warmup 10 epochs + Cosine decay
- MixUp (α=0.8) + RandAugment + Label Smoothing (0.1)
- Gradient Clipping (max_norm=1.0)

**训练结果示例：**

```
Epoch 500/500  19s
  Train loss=1.1610 acc=87.16%
  Val   loss=0.3222 acc=94.54%
```

### 4. MAE 预训练（Masked Autoencoder）

采用 [He et al. (2022)](https://arxiv.org/abs/2111.06377) 的 MAE 范式，随机遮盖图像 patch 并重建进行自监督预训练。

```bash
python src/vit/vit_cifar10_mae_pretrain_stax.py
```

- **Encoder**：7 层 ViT（无 CLS），仅处理可见 patch（~40%）
- **Decoder**：2 层轻量 Transformer（192d），处理完整序列
- **掩码率**：60%
- **目标**：MSE 重建被遮盖 patch 的像素值
- 100 epochs, batch_size=128, AdamW (lr=1.5e-4)

### 5. MAE 微调

加载预训练 Encoder，添加 CLS token + 分类头，监督微调。

```bash
python src/vit/vit_cifar10_mae_finetune_stax.py
```

- 200 epochs, batch_size=128, AdamW (lr=1e-4)
- MixUp + RandAugment + Label Smoothing

## 断点恢复

`vit_cifar10_stax.py`、`swin_cifar10_stax.py` 和 `swin_cifar10_cnn_stem.py` 支持自动 checkpoint 保存/恢复。训练中断后再次运行会自动加载最近的 checkpoint，从中断的 epoch 继续训练，无需手动干预。

| 脚本 | Checkpoint |
|------|------------|
| ViT 监督 | `ckpt/vit_cifar10_supervised.pkl` |
| Swin 监督 | `ckpt/swin_cifar10_supervised.pkl` |
| Swin + CNN PreStem 监督 | `ckpt/swin_cnn_stem_cifar10_supervised.pkl` |
| MAE 预训练 | `ckpt/vit_cifar10_mae_pretrain.pkl` |
| MAE 微调 | `ckpt/vit_cifar10_mae_finetune.pkl` |

## 核心技术细节

### Flash Attention — 自定义 VJP

两个版本的 IO-aware Flash Attention，均包含手写前向/反向 kernel：

| 特性 | v1 | v2 |
|------|-----|-----|
| 前向循环 | Q-outer, KV-inner | Q-outer, KV-inner |
| 反向循环 | Q-outer, KV-inner | **KV-outer, Q-inner** |
| 累计方式 | m + ℓ 分两变量 | Õ (unscaled) 统一累加 |
| 参考 | Dao 2022 | Dao 2023 |

已修复分块模式下 epsilon 不一致导致的梯度漂移问题：前向归一化和 log-sum-exp 保存统一使用 `_MIN_NORMALIZER = 1e-12`。

### Stax 扩展库 (`stax_plus.py`)

基于 JAX 官方 Stax 示例增强：
- `BatchNorm` — EMA 动量，PyTorch 对齐，JIT 兼容
- `LayerNorm` — PyTorch 等价实现
- `GroupNorm` / `Dropout` / `Lambda`
- `serial` / `parallel` 组合器（自动线程 BN state）
- 便利函数 `dropout()` — 供手动 apply 使用

### Transformer 模块

`transformer_flash_cudnn.py` (v1) 和 `transformer_flash_v2.py` (v2) 提供：
`MultiHeadSelfAttention` / `MultiHeadCrossAttention` / `TransformerEncoderBlock` /
`TransformerDecoderBlock` / `TransformerEncoder` / `TransformerDecoder` / `Transformer`

所有层遵循 Stax `(init_fun, apply_fun)` 约定，可直接导入复用。

### 训练技巧

| 技巧 | 说明 |
|------|------|
| MixUp | α=0.8，温和混合样本与标签 |
| RandAugment | N=2, M=9 自动增强 |
| Label Smoothing | 0.1，防止过拟合 |
| Warmup + Cosine Decay | 线性 warmup → 余弦衰减到峰值 1% |
| Gradient Clipping | max_norm=1.0 |
| AdamW | 解耦权重衰减 |

## 依赖

```
jax[cuda12]          # JAX + CUDA 支持
optax                # JAX 优化器
torchvision / torch  # CIFAR-10 数据加载
matplotlib           # 训练曲线绘制
numpy                # 数值计算
```

## 运行环境说明

- `XLA_FLAGS=--xla_gpu_enable_command_buffer=` — 禁用 CUDA 命令缓冲区（兼容 CUDA 13.x 驱动 + CUDA 12.x 运行时）
- PyTorch DataLoader 使用 `num_workers=0`（避免与 JAX CUDA 线程模型冲突）
- ViT 图像转置：PyTorch `(C, H, W)` → JAX Stax 期望的 `(H, W, C)`
- Swin 图像格式：`(B, C, H, W)`（PyTorch 兼容）

## 训练流程示例

```bash
# 监督训练（ViT）
python src/vit/vit_cifar10_stax.py

# 监督训练（Swin）
python src/vit/swin_cifar10_stax.py

# 监督训练（Swin + CNN PreStem）
python src/vit/swin_cifar10_cnn_stem.py

# MAE 两阶段训练
python src/vit/vit_cifar10_mae_pretrain_stax.py    # → ckpt/vit_cifar10_mae_pretrain.pkl
python src/vit/vit_cifar10_mae_finetune_stax.py     # 加载预训练权重
```

## 输出

每个训练脚本会生成：
- `*.npz` — 训练曲线数据（loss、accuracy）
- `*.png` — 可视化训练曲线
- `ckpt/*.pkl` — 模型权重 / 断点恢复文件

## 参考资料

- [An Image is Worth 16x16 Words](https://arxiv.org/abs/2010.11929) — ViT
- [Masked Autoencoders Are Scalable Vision Learners](https://arxiv.org/abs/2111.06377) — MAE
- [FlashAttention: Fast and Memory-Efficient Exact Attention](https://arxiv.org/abs/2205.14135) — FlashAttention v1
- [FlashAttention-2: Faster Attention with Better Parallelism](https://arxiv.org/abs/2307.08691) — FlashAttention-2
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030) — Swin

## License

MIT