# 一个vulkan写的AI程序

**Repository Path**: maple1618/vkAI

## Basic Information

- **Project Name**: 一个vulkan写的AI程序
- **Description**: 一个vulkan写的AI程序，用来测试vulkan是否能够跑通一个完整AI任务
- **Primary Language**: C++
- **License**: LGPL-2.1
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-10-14
- **Last Updated**: 2026-10-19

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# learn04 — ConvSIREN 单图 INR（Vulkan compute）

用一组 compute shader 把一个 512×512 / 1920×1080 的 RGB 图拟合进一个 SIREN 风格的卷积隐式表示，
训练与推理共用同一份模型结构。全流程只有 Vulkan + Slang，没有 PyTorch 依赖（Python 只用来做参考对拍）。

## 1. 构建模型（两阶段）

```
slang/*.slang ──bulid(cook)──> res_train/pass*.spv ──embed.py──> resources.S/.cpp
                                                              └──> 运行期二进制（不含 Slang 编译器）
```

- `bulid` 只做编译（Slang → SPIR-V），产物落构建目录，源码树零产物
- 运行期二进制把 SPV 内嵌为二进制资源，链接时不再依赖 Slang
- main 由 `-DMAIN_SRC=` 选择：`conv_train.cpp`（训练相）/ `conv_run.cpp`（生产相）

## 2. 训练

最快路径（自动建 `build_<图名>/`、强制按该图重烘 SPV、可断点续训）：

```bash
./train_kodak.sh /path/to/image.png 6000
```

手动路径：

```bash
source /maple/share/vulkan/setup-env.sh        # Vulkan SDK + Slang

ACER_TARGET_PNG=/path/to/image.png \
cmake -S . -B build_demo -DMAIN_SRC=conv_train.cpp -DACER_TRAIN=ON
ACER_TARGET_PNG=/path/to/image.png cmake --build build_demo -j"$(nproc)"

./build_demo/conv_train 6000 /path/to/image.png build_demo/runs/demo
```

> `ACER_TARGET_PNG` 是给 **cook** 用的：着色器里的几何常量（`H2/W2/PAD/IMG/IMGW`）在编译期
> 烘进 SPIR-V。换图必须换构建目录并重新 cook，否则运行期 buffer 与 shader 常量错位。

命令行参数顺序无关：`conv_train [轮数] [png 路径] [run 目录]`，轮数是**本次增量**，
续训直接再跑一次同一条命令。

### 产物布局

```
build_<图名>/
  res_train/               该图的 SPV（cook 产物）
  model/                   交付模型包：model.cpp + model.flags + model.spvN + model_s0_b0.bin + model.geom
  runs/<图名>/
    1-0.bin                续训权重
    58-2.bin 58-3.bin      Adam 一阶/二阶矩
    58-4.bin               训练配置（轮数/lr/EMA 等）
    59-1.bin               Muon 动量
    loss.csv               每轮 mse / psnr / lr / wall_ms（追加写）
    out_<轮数>.png         收尾重建图（肉眼核对）
```

## 3. 推理

模型包是自包含的（SPV + 权重 + 几何），单独建一个构建目录读取即可，不必切回训练目录：

```bash
source /maple/share/vulkan/setup-env.sh        # Vulkan SDK + Slang (configure 期也要)
cmake -S . -B build_run -DMAIN_SRC=conv_run.cpp \
      -DMODEL_DIR=$PWD/build_demo/model
cmake --build build_run -j"$(nproc)"

cd build_run && ./conv_run /path/to/image.png   # 给了目标图就顺带打印 PSNR
# → reconstructed_gpu.png（尺寸由 model.geom 决定，写在当前目录）
```

`conv_run` 的 cook 相是空转：19 趟 SPV 随模型包（`model.cpp` 的 `.incbin`）一起链进来，
所以 `build_run/res` 里没有 `pass*.spv` 是正常的。

`model.geom` 记录 `IMG_W IMG_H TILE_W TILE_H PAD`，推理端据此逐 tile 提交，
因此训练用的任何分辨率/tile 切分都能原样复现。

> ★ 模型包里的权重是 `.incbin` **链接期**嵌进二进制的（`model.cpp`）。训练端重新
> `exportModel` 之后必须再 `cmake --build build_run`（即使源码没改），否则跑的还是
> 上一次链进去的旧权重。

## 4. 几何与显存

- 感受野半径 `PAD = 15`（1+2+4+8−1，对应 4 层 `dilation=1,2,4,8`）
- 运行期从目标图探测尺寸；整图能放进 `VRAM_BUDGET = 9 GiB` 时用 `PAD=0` 单 tile（零 halo）
- 否则在预算内搜索 halos 白算最少的**二维 tile**（比一维分带重叠更少）：
  1920×1080 / NF=512 时选 `480×540`、8 个 tile，halo 约 12%
- `NF = 512`（`conv_train.cpp` 顶部常量），改宽度后 shader 端常量由 `writeSlang` 自动注入

## 5. 损失与优化器

- `loss = MSE + CHROMA_LAMBDA·chroma + GRAD_LAMBDA·grad + SSIM_WEIGHT·(1−SSIM)`
- 每 tile 只在自己的 core 区内算差分/SSIM；`mse_sum` 把部分和按 tile 切片写回，host 用 double 累加
- Adam（UPD 趟）+ Muon（GEMM 参数，牛顿-舒尔茨 5 步）+ 末段权重平均（`avgN` 轮）
- lr 按**绝对轮号** warmup/平台衰减，续训不会重新 warmup

## 6. 参考对拍（可选）

`train.py` / `conv_ref.py` / `conv_grad_ref.cpp` 是 PyTorch 参考实现，用来逐趟比对梯度
（注意布局：GPU 侧激活是 HWC，torch 参考是 NCHW；只有 `conv_ref_out/target/dlout.bin` 是 HWC）。

## 7. 目录说明

| 路径 | 内容 |
|---|---|
| `conv_train.cpp` / `conv_run.cpp` | 训练相 / 生产相（当前主线） |
| `slang/conv/` | 训练/推理用到的全部着色器 |
| `vulkan/` | 框架层（`acer`：buffer/命令/shader 容器、模型导出/加载） |
| `mini/` `iris/` `ptbxl/` `slang/iris` `slang/ptbxl` | 其他实验，独立于 conv 主线 |
| `bench_*.cpp` | dw 算子微基准（独立编译，不参与主目标） |
| `TRAINING_NOTES.md` | 训练运行记录（含作者机器上的绝对路径，仅作历史） |

## 8. Debug 构建的 sanitizer

默认 `-DCMAKE_BUILD_TYPE=Debug` 带 ASan/UBSan。`conv_train` 里用 LSan 官方钩子屏蔽了
NVIDIA 驱动 dlopen 的 libdbus 外部泄漏，其余泄漏仍会照常报告；要连屏蔽统计一起静默：

```bash
ASAN_OPTIONS=print_suppressions=0 ./build_demo/conv_train ...
```