diff --git a/.claude/skills/train-backbones/SKILL.md b/.claude/skills/train-backbones/SKILL.md new file mode 100644 index 0000000..8eb2f1e --- /dev/null +++ b/.claude/skills/train-backbones/SKILL.md @@ -0,0 +1,121 @@ +--- +name: train-backbones +description: Launch the BoardOCR backbone sweep. Asks the user for launch mode, image size, epoch count, and optional backbone subset; auto-tunes batch_size/num_workers/prefetch_factor from detected GPU count + VRAM + CPU cores via scripts/autotune.py; then runs scripts/train_backbones.sh with the resolved env vars. Use for the "start / kick off / run the training sweep" ask on this project. +--- + +# train-backbones + +Kicks off training of BoardOCR across one or more backbones. Handles both the +1-GPU dev box (RTX 4070 Ti) and the 8-GPU A100 target uniformly by delegating +to `scripts/train_backbones.sh` (which supports three launch modes) and letting +`scripts/autotune.py` size DataLoader/batch based on the detected hardware. + +## What this skill does + +1. Ask the user (via `AskUserQuestion`) for the choices they should actually + own — launch mode, image size, epochs, optionally the backbone subset. Do + NOT ask for anything that autotune can infer (workers, batch, VRAM, etc.). +2. Detect hardware and derived hyperparams by running + `uv run python scripts/autotune.py --backbone --image-size --mode [--nprocs ]`. + Parse its `KEY=VALUE` lines. The backbone passed to autotune should be the + *heaviest* one in the sweep, so batch_size holds for every backbone in the run. +3. Print a compact plan back to the user (mode, backbones, image_size, epochs, + resolved batch/workers) and then launch the sweep by exporting the env vars + inline with the shell invocation: + `EPOCHS=... BATCH_SIZE=... NUM_WORKERS=... PREFETCH_FACTOR=... IMAGE_SIZE=... ./scripts/train_backbones.sh` +4. Stream stdout so the user sees the sweep script's progress log. + +## The three launch modes + +Present them via a single `AskUserQuestion` with these labels/descriptions so +the user picks based on their goal, not implementation: + +- **Sequential (1 backbone × 1 GPU)** — simplest, dev-box friendly. Runs each + backbone in turn. Env: default. +- **DDP (1 backbone × N GPUs)** — one heavy backbone, split across all GPUs + via torchrun. Env: `NPROC_PER_NODE=`. Useful for training a single + chosen model as fast as possible. +- **Parallel (N backbones × 1 GPU each)** — best for a full comparison sweep + on a multi-GPU node. Each backbone runs on its own GPU concurrently. Env: + `PARALLEL_GPU=1`. Waves of `NGPUS` at a time. + +If autotune reports `NGPUS_DETECTED=1`, only Sequential makes sense — skip the +question and note in the plan why. + +## Image size options + +- **224** (default) — matches ImageNet pretraining, smallest cache (~10 GB), + fastest per epoch. Each 9x9 board cell gets ~25 px. +- **288** — ~32 px/cell, moderate cache (~17 GB), ~1.5x epoch time. +- **384** — ~42 px/cell, big cache (~31 GB), ~2.5x epoch time. Best for + distinguishing fine-detail characters like 成香/成桂. + +## Epoch options + +Present 30 / 50 / 100 / 200 as anchor choices. Note that first epoch of a +cold-cache run pays the preload build cost (~30-60 s for full dataset). + +## Backbones + +Default = all 9 (small → large): `mobilenet_v3_small mobilenet_v3_large +convnext_atto convnext_femto efficientnet_b0 convnext_pico efficientnet_b1 +convnext_nano convnext_tiny`. Ask only if the user hints at a subset. Pass via +the `BACKBONES` env var, space-separated. + +## VRAM cache — not implemented + +The user may ask about loading the preload cache into VRAM (an A100 can hold +the whole 10 GB cache easily). This is a known followup, NOT implemented today. +Reasons: + +- The current pipeline decodes preload → DRAM → CPU DataLoader workers → GPU. +- Moving the cache to VRAM only pays off if augmentation also runs on GPU, + because otherwise workers still have to copy back to CPU per sample. +- That requires porting the Albumentations chain to Kornia — a real refactor, + not a one-line change. + +If asked, explain the tradeoff honestly and defer. The DRAM cache is already +near-instant on cache hit (mmap) and DataLoader isn't currently the bottleneck. + +## Autotune calling convention + +`scripts/autotune.py` prints eval-able env assignments. Call it once with the +*heaviest* backbone in the sweep so batch_size is safe for every backbone +(smaller ones will fit trivially). + +Examples: +```bash +# Parallel mode on the 8-GPU node, all backbones: +uv run python scripts/autotune.py --backbone convnext_tiny --image-size 224 --mode parallel + +# DDP mode, 8 GPUs, single backbone: +uv run python scripts/autotune.py --backbone convnext_tiny --image-size 224 --mode ddp --nprocs 8 +``` + +## Full launch invocation + +Compose the final shell command. Mode → env var mapping: +- sequential: no extra env +- ddp: `NPROC_PER_NODE=${NGPUS_DETECTED}` +- parallel: `PARALLEL_GPU=1` + +Then run (adjust the trailing script name if the user wants to change it): + +```bash +IMAGE_SIZE= EPOCHS= BATCH_SIZE= NUM_WORKERS= PREFETCH_FACTOR=

\ + [BACKBONES=""] [PARALLEL_GPU=1 | NPROC_PER_NODE=] \ + ./scripts/train_backbones.sh +``` + +Show it in the plan so the user can copy/rerun without going through the +skill later. + +## Things not to ask about (skill owns these) + +- Learning rate — leave the script default (3e-4). If asked, mention linear + scaling for larger effective batch and that the user can pass `LR=<...>` env. +- Preload — always on (`--preload`). Cache is on-disk mmap; zero cost after + first build. +- Checkpoint dir / W&B project name — script defaults are fine. +- `RESUME_INCOMPLETE` — leave off unless the user asks to resume; not a + hyperparameter question. diff --git a/.devcontainer/compose.yaml b/.devcontainer/compose.yaml index 4277966..b19bcd4 100644 --- a/.devcontainer/compose.yaml +++ b/.devcontainer/compose.yaml @@ -6,10 +6,14 @@ services: volumes: - ../:/home/vscode/app:cached - venv:/home/vscode/app/.venv - shm_size: 4gb + shm_size: 32gb tty: true stdin_open: true volumes: venv: driver: local + hf-cache: + driver: local + uv-cache: + driver: local \ No newline at end of file diff --git a/.devcontainer/cuda/Dockerfile b/.devcontainer/cuda/Dockerfile index d930d54..2006da7 100644 --- a/.devcontainer/cuda/Dockerfile +++ b/.devcontainer/cuda/Dockerfile @@ -1,11 +1,33 @@ FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 +ARG USERNAME=vscode +ARG USER_UID=1000 +ARG USER_GID=$USER_UID + +# Keep BuildKit's apt cache mounts populated across builds: docker-clean wipes +# the archives after every install, so drop it and tell apt to keep them. +RUN rm -f /etc/apt/apt.conf.d/docker-clean \ + && echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' \ + > /etc/apt/apt.conf.d/keep-cache + # NVIDIA base image is minimal; install the tools required by common-utils and # subsequent devcontainer features. Everything else (uv, zsh, gh, claude-code) # is layered on via features in devcontainer.json. -RUN apt-get update && apt-get install -y --no-install-recommends \ +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ + DEBIAN_FRONTEND=noninteractive apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ ca-certificates curl git sudo openssh-client tzdata locales \ - && locale-gen en_US.UTF-8 \ - && rm -rf /var/lib/apt/lists/* + && locale-gen en_US.UTF-8 ENV LANG=en_US.UTF-8 LC_ALL=en_US.UTF-8 + +# Create the non-root vscode user the devcontainer runs as (remoteUser). The +# NVIDIA base is root-only, so unlike the devcontainers base image we make it +# here with passwordless sudo. +RUN groupadd --gid $USER_GID $USERNAME \ + && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME -s /bin/bash \ + && echo "$USERNAME ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/$USERNAME \ + && chmod 0440 /etc/sudoers.d/$USERNAME + +USER $USERNAME diff --git a/.devcontainer/cuda/compose.yaml b/.devcontainer/cuda/compose.yaml index 4493240..40406cb 100644 --- a/.devcontainer/cuda/compose.yaml +++ b/.devcontainer/cuda/compose.yaml @@ -8,9 +8,10 @@ services: # workspace root is two levels up. - ../../:/home/vscode/app:cached - venv:/home/vscode/app/.venv - # Persistent HF datasets cache so `HFCaptureDataset` doesn't re-download - # the ~21GB parquet shards on container rebuild. - - hf-cache:/home/vscode/.cache/huggingface + # Persistent per-user cache dir. Covers HF datasets (~21GB parquet + # shards for `HFCaptureDataset`), uv package cache, and anything else + # that would otherwise cost time to rebuild on container recreation. + - cache:/home/vscode/.cache # DataLoader workers on CUDA can consume more shared memory than the # default 64MB; 16GB matches typical CUDA training rigs. shm_size: 16gb @@ -29,5 +30,5 @@ services: volumes: venv: driver: local - hf-cache: + cache: driver: local diff --git a/.devcontainer/cuda/devcontainer-lock.json b/.devcontainer/cuda/devcontainer-lock.json new file mode 100644 index 0000000..1b4b3bb --- /dev/null +++ b/.devcontainer/cuda/devcontainer-lock.json @@ -0,0 +1,39 @@ +{ + "features": { + "ghcr.io/devcontainers-extra/features/fzf:1": { + "version": "1.0.15", + "resolved": "ghcr.io/devcontainers-extra/features/fzf@sha256:dbac92d89862c0f266453772ea0e73089d520ff66544f9c146480677fc2dbf7d", + "integrity": "sha256:dbac92d89862c0f266453772ea0e73089d520ff66544f9c146480677fc2dbf7d" + }, + "ghcr.io/devcontainers/features/common-utils:2": { + "version": "2.5.9", + "resolved": "ghcr.io/devcontainers/features/common-utils@sha256:cb0c4d3c276f157eed17935747e364178d75fee17f55c4e129966f64633deb3a", + "integrity": "sha256:cb0c4d3c276f157eed17935747e364178d75fee17f55c4e129966f64633deb3a" + }, + "ghcr.io/devcontainers/features/docker-outside-of-docker:1": { + "version": "1.10.0", + "resolved": "ghcr.io/devcontainers/features/docker-outside-of-docker@sha256:c2c2cf829505ead8e4892c88c31b6594ae94a2bbb209e16e1fac456c1a3a624e", + "integrity": "sha256:c2c2cf829505ead8e4892c88c31b6594ae94a2bbb209e16e1fac456c1a3a624e" + }, + "ghcr.io/devcontainers/features/github-cli:1": { + "version": "1.1.0", + "resolved": "ghcr.io/devcontainers/features/github-cli@sha256:d22f50b70ed75339b4eed1ba9ecde3a1791f90e88d37936517e3bace0bbad671", + "integrity": "sha256:d22f50b70ed75339b4eed1ba9ecde3a1791f90e88d37936517e3bace0bbad671" + }, + "ghcr.io/devcontainers/features/python:1": { + "version": "1.8.0", + "resolved": "ghcr.io/devcontainers/features/python@sha256:fbcad6955caeecc5ad3f7886baf652e25cba5225a6c4c2287c536de2e5607511", + "integrity": "sha256:fbcad6955caeecc5ad3f7886baf652e25cba5225a6c4c2287c536de2e5607511" + }, + "ghcr.io/jsburckhardt/devcontainer-features/uv:1": { + "version": "1.0.0", + "resolved": "ghcr.io/jsburckhardt/devcontainer-features/uv@sha256:542a0bc2203205b3c696de650ba862f280b20af3543493cc232edd9ac35791f7", + "integrity": "sha256:542a0bc2203205b3c696de650ba862f280b20af3543493cc232edd9ac35791f7" + }, + "ghcr.io/stu-bell/devcontainer-features/claude-code:0": { + "version": "0.1.0", + "resolved": "ghcr.io/stu-bell/devcontainer-features/claude-code@sha256:f87b4da3f8648db9111cb25c6bd3817d096018d81f3fb596340f7ffdba14409a", + "integrity": "sha256:f87b4da3f8648db9111cb25c6bd3817d096018d81f3fb596340f7ffdba14409a" + } + } +} diff --git a/.devcontainer/postCreateCommand.sh b/.devcontainer/postCreateCommand.sh index 200dbfc..8d9c429 100644 --- a/.devcontainer/postCreateCommand.sh +++ b/.devcontainer/postCreateCommand.sh @@ -1,4 +1,5 @@ #!/bin/sh -sudo chown -R "$(whoami)":"$(whoami)" /home/"$(whoami)"/app/.venv -uv sync +sudo chown -R $(whoami):$(whoami) /home/$(whoami)/app/.venv +sudo chown -R $(whoami):$(whoami) ~/.cache +uv sync --extra experiment diff --git a/.gitignore b/.gitignore index f3ad4d3..3913635 100644 --- a/.gitignore +++ b/.gitignore @@ -259,6 +259,7 @@ data/ # 学習中間物 runs/ wandb/ +logs/ checkpoints/ *.pt *.pth diff --git a/TRAINING_PLAN.md b/TRAINING_PLAN.md index 07a4a2e..790cf7c 100644 --- a/TRAINING_PLAN.md +++ b/TRAINING_PLAN.md @@ -3,29 +3,33 @@ board_ocr の学習を回しながら気づいた改善候補と、次のイテレーションの優先度メモ。 `docs/` は MITO 側から同期する仕様書用なので、内部計画はここに置く。 -## 現状の観測（2026-07-08 board_ocr full run） +## 現状の観測（2026-07-12 backbone sweep） -前回 20 エポック → 追加 40 エポック（`RESUME=latest EPOCHS=60`）で継続学習中。 +`scripts/train_backbones.sh` で 9 backbone × 50 epoch × image_size=224 の sweep 実施済み。 +詳細な数字と考察は [`docs/ocr-scaling-outlook.md`](./docs/ocr-scaling-outlook.md) に集約。 -分布可視化ツール：`scripts/inspect/analyze_hand_distribution.py`（出力: `runs/hand-distribution.png`） +sweep のサマリ（val, epoch 50）： -| epoch | loss | board_loss | hand_loss | cell_acc | sfen_acc | -|------:|-----:|-----------:|----------:|---------:|---------:| -| 5 | 0.329 | 0.025 | 0.607 | 0.993 | 0.046 | -| 20 | 0.120 | 0.007 | 0.226 | 0.998 | 0.324 | -| 30 | 0.099 | 0.006 | 0.187 | 0.998 | 0.417 | +| Backbone | Params | val cell_acc | val sfen_full | +|---|---:|---:|---:| +| mobilenet_v3_small | 1.1M | 0.966 | 0.385 | +| mobilenet_v3_large | 3.3M | 0.989 | 0.593 | +| convnext_atto | 3.5M | 0.991 | 0.611 | +| efficientnet_b0 | 4.4M | 0.991 | 0.661 | +| convnext_femto | 4.9M | 0.992 | 0.645 | +| efficientnet_b1 | 6.9M | 0.992 | 0.663 | +| convnext_pico | 8.7M | 0.992 | 0.634 | +| convnext_nano | 15.1M | 0.993 | 0.660 | +| convnext_tiny | 28.0M | 0.994 | 0.716 | -epoch 20 時点の val 側： -- `val/board/cell_acc = 0.999`（train より高い＝過学習なし） -- `val/hand/slot_acc = 0.939` -- `val/hand/full_acc = 0.385` -- `val/sfen/full_acc = 0.369` +分布可視化ツール：`scripts/inspect/analyze_hand_distribution.py`（出力: `runs/hand-distribution.png`） ### 読み取れる状態 -- **board head は epoch 5 前後で飽和**（cell_acc 0.99+ で伸びしろほぼ無し）。 -- **hand head はまだ学習中**（hand_loss / sfen_acc とも右肩下がり継続）。ボトルネックはこちら。 -- val > train の傾向で過学習は現状ゼロ。エポックはさらに伸ばせる。 +- **board head は epoch 5〜10 前後で飽和**（cell_acc 0.99+ で伸びしろほぼ無し）。全 backbone 共通。 +- **hand head がボトルネック**。sfen_full は 0.385〜0.716 と 33pt 開いており、backbone 容量と学習余地の両方が効いている。 +- val > train の傾向で過学習は現状ゼロ。エポックはさらに伸ばせるが、後半勾配は落ち始めている backbone も多い。 +- **参考: 旧 v1/v2 の観測**（image_size=288, batch=128, lr=6e-4）は本ファイル末尾の「v1 baseline」「v2」節に残す。sweep は条件が違うので直接比較不可。 ## 改善候補 @@ -48,18 +52,35 @@ epoch 20 時点の val 側： ### B. データ分布の偏り対応（優先度：高） -**観測**（`scripts/inspect/analyze_hand_distribution.py`、`data/train.jsonl` n=27,000 局面 = 378,000 スロットラベル）： - -- ラベル `0`（=空スロット）が **55.94%**。`1` が 26.09%、`2` が 8.65%、以降 3%以下まで急落。 -- モデルが常に「0」と答えても素の accuracy が 56% 出せる構造。 -- 駒種別の非ゼロ率： - - 歩：**90%以上**（唯一 10枚超まで散らばる） - - 金・銀・桂：43〜46% - - 香・角：30〜34% - - **飛：19〜20%**（ほぼ 0 or 1 枚） +**観測**（`scripts/inspect/analyze_hand_distribution.py` を現行 `ultemica/piyoshogi` `ocr_paired` train n=18,000 SFEN = 252,000 スロットラベルに対して 2026-07-12 実測）： + +- ラベル `0`（=空スロット）が **65.68%**。`1` が 21.20%、`2` が 6.39%、`3` 以降は 3% 未満に急落し、`10` は 0.12%、`18` に至っては 0.003%。 +- モデルが常に「0」と答えても素の accuracy が **65.68%** 出せる構造。旧集計（55.94%）より上振れしており、0 バイアスは寧ろ強まっている。 +- 駒種別の非ゼロ率（先手/後手はほぼ対称）： + - 歩（S:P / G:P）：**68.03% / 63.95%**（唯一 count=1〜18 全域に散らばる。count=10 も 148/145 件、count=18 も 4/4 件） + - 角（S:B / G:B）：42.94% / 41.51% + - 金（S:G / G:G）：32.48% / 32.34% + - 銀（S:S / G:S）：31.42% / 30.69% + - 桂（S:N / G:N）：29.28% / 28.84% + - 香（S:L / G:L）：22.24% / 21.79% + - **飛（S:R / G:R）：17.84% / 17.13%**（ほぼ 0 or 1 枚、稀に 2 枚） - 同じ hand head で 14 スロット一括処理しているが、実質**駒種ごとに難易度が全く違うタスク**。 +- 物理制約（飛/角は最大 2、香桂銀金は最大 4、歩のみ 18 まで）はデータ上も遵守されており、`PIECE_MAX_PER_SLOT` の logit マスクは正しく効いている。 + +**現行データでの raw class weight（`grand_total / (HAND_MAX_COUNT * v)`）**： +count=0 → 0.080、count=6 → 10.27、count=10 → 45.3、count=15 → 221、count=18 → 1658。5〜6 桁のレンジ。 + +**val 分布の穴（2026-07-12 実測, n=2,000 SFEN）**： +val は train と分布形が違う。 +- ラベル 0 の比率が **75.10%**（train 65.68%）。「常に 0」ベースラインの val slot_acc が **10pt 底上げ**される。 +- **count=11 以上が val に 1 件も無い**。歩 count=11〜18 は train には計 500 件超あるが val では見えないため、**regression head の高枚数改善効果を val slot_acc で評価できない**。 +- count=10 は val 全体で 1 件のみ（S:P）。 +- 飛の非ゼロ率は val で **9.4% / 7.2%**（train は 17.8% / 17.1%）。val でさらに希少。 +- 香/桂/銀/金は val 側でスロット別の max が偶発的に非対称（例：S:L max=2、G:L max=4）。per-slot accuracy の左右比較に要注意。 + +→ **hand の改善評価は `val/hand/slot_acc` だけ見ない**。per-count recall（count ≥ 3）や、高枚数用の合成 mini-eval セット、`test-realistic` 側の実測を併用する。 -**素朴な `1 / freq` の危険性**：クラス重みが 0.094（count=0）〜 9947（count=16）まで**5 桁開く**。希少クラスの勾配が暴発して学習不安定になるので、そのままは NG。 +**素朴な `1 / freq` の危険性**：クラス重みが 5〜6 桁開く（count=0 で 0.08、count=18 で 1658）。希少クラスの勾配が暴発して学習不安定になるので、そのままは NG。 **対応候補**： @@ -146,29 +167,16 @@ epoch を長めに指定して回すぶん、val 指標が K エポック改善 ## 既知のデータ制約 -### 持ち駒 count=10 のギャップ（要対応） +### 持ち駒 count=10 のギャップ（解消済み） -`data/train.jsonl` の全 27,000 局面を集計しても、**どのスロットにも count=10 が 1 件も現れない**（count=9 は 585 件、count=11 は 158 件あるのに）。加えて count=17, 18 も完全に不在。 +以前は `data/train.jsonl` に **どのスロットにも count=10 が 1 件も現れない**問題があった（piyo-hook 側の SFEN 出力バグに起因）。特に歩スロット（S:P, G:P）は実局面で 10 枚以上が普通に発生するため、学習データの穴が推論時に「10 枚 → 8/9/11 に丸める」挙動を起こしていた。 -**原因**：ぴよ将棋（データ生成元）側で 10 枚持ちの局面が SFEN 出力の段階で生成されない既知のバグ。mito-train 側で修正できるものではなく、入力データの制約として顕在化している。 - -**⚠️ 歩は要注意**：**実局面では歩 10 枚（およびそれ以上）の局面は普通に発生する**。飛・角・金・銀・桂・香は 10 枚以上を持てないので count=10 のギャップは実害ゼロだが、**歩のスロット（S:P, G:P）だけは実運用推論で count=10 を要求される場面がある**。学習データに 10 枚が 1 件も無い状態だと、モデルは歩 10 枚を 9 or 11 に誤読する挙動になる。 - -**推論への影響**： - -- 歩以外のスロット：学習時に count=10 が来ないだけでなく実運用でも来ないので影響なし。 -- **歩スロット**：実局面で count=10 が来たとき、そのラベルの logit は全く学習されていない → 8/9/11 に丸める挙動。 -- count=17, 18 は理論上歩でのみあり得るが実局面での出現頻度は極めて低い。優先度は 10 より低い。 - -**対策候補**： +**現状**：`ultemica/piyoshogi` の最新スナップショットには歩 count=10 の局面が含まれており、この構造的な穴は解消済み。歩の高枚数（10 以上）は依然として希少（分布としては裾）だが、logit が「完全ゼロ学習」状態ではなくなっている。 -- **A. 歩スロットだけデータ合成で 10 枚局面を追加**（本命）。 - - piyo-hook 修正を待つ間の暫定策として、既存の 9 枚 / 11 枚局面から歩コマ画像を貼り替えて 10 枚版を合成。 - - 数百件でも「count=10 の logit が完全ゼロ学習」状態は解消できる。 - - 対象を歩スロットに限定すれば副作用が少ない。 -- **B. piyo-hook 側のバグ修正を issue 化**（本筋）。上流で直せば A は不要になる。 +**残っている hand 側の課題**（詳細は [`docs/ocr-scaling-outlook.md`](./docs/ocr-scaling-outlook.md) の「hand の失敗パターン」を参照）： -優先度：**A を短期で、B を並行で issue 化**。C（駒種別の出力次元制約）は下記の別項で扱う。 +1. **枚数の隣接ミス**（本命）: 19-way CE のため「5 vs 6」も「5 vs 18」も同じ loss。序数情報が捨てられている。→ 対策は本ファイル **§A の regression head**。 +2. **0 バイアス**: ラベル 0（空スロット）が過半、飛スロットは 0/1 でほぼ完結。希少枚数（3〜9）で under-count しやすい。→ v2 で入れた **§B の class weight sqrt+clip** が対症療法。本丸は §A。 ### 駒種別の理論最大枚数（モデル設計に反映すべき制約） @@ -229,3 +237,124 @@ epoch を長めに指定して回すぶん、val 指標が K エポック改善 **config**：`EPOCHS=60 BATCH_SIZE=128 LR=6e-4 IMAGE_SIZE=288 BACKBONE=mobilenet_v3_small` **ckpt出力**：`runs/board-ocr-v2/` **目標**：epoch 60 で sfen_acc > 0.55（v1 の 45 epoch 相当を超える） + +## v3（HF natural + synthetic マージ版） + +### 見つけた偏り + +v2 は HF `ultemica/piyoshogi` `ocr_paired` train（18,000 SFEN、`build_manifests.py` で作った natural 40% + opening 20% + synthetic 40% + existing のミックス）を教師に使っていた。そこには **hand 側に強い分布バイアス** があった（2026-07-12 実測）： + +| 指標 | v2 データ (HF 18k) | +|---|---:| +| slot ラベル 0（空スロット）の比率 | **65.68%** | +| slot ラベル 1 の比率 | 21.20% | +| slot ラベル 3 以上の合計 | ~4% | +| 歩 count=7 以上の合計 | ほぼゼロ | +| 歩 count=18 の観測 | 4 件（train 全体で） | + +- 常に 0 を答えるだけで slot_acc **65.68%** が出る構造。hand head の勾配が「0 を当てにいく」方向に寄る。 +- 高枚数（3〜18）は裾でしか観測されず、隣接ミス（5 vs 6, 10 vs 11）の学習信号が薄い。 +- 飛 (R) / 角 (B) は非ゼロ率が **17〜42%** に留まる。 +- val 側（2,000 SFEN）は train より更に極端で、**count=11 以上が 1 件も無い**。「regression head を入れたい」って言っても val で効果測定ができない。 + +詳細は `docs/ocr-scaling-outlook.md#hand-の失敗パターン` に集約。 + +### 偏りの修正手順 + +**方針**：HF 18k は捨てず、そこに **裾を厚くした synthetic 10k を足す**。ただし単純ユニオンだと HF 側の 0 バイアスがそのまま重みで持ち込まれるので、**HF 側も rare-preserving で 10k に間引き、10k + 10k = 20k の均等 mix** に落ち着かせる。synth 側の SFEN は piyo-hook で HF と同じ 4 端末（iPhone10,1 / 11,8 / 15,4 / iPad14,10）で撮影して、**画像スタイルは 20k 全件で統一**する。合成レンダラは使わない（画像スタイルと分布の相関を学習するリスクを避ける）。 + +1. **`scripts/data/generate_synthetic_sfens.py`**: `python-shogi` でランダム対局を回し、途中局面から **piece type ごとに Uniform[0, PIECE_MAX] で hand をサンプリング**。盤上の駒を hand に移すだけで生成するので、nifu / per-type 総数保存 / logit マスクとの整合はすべて自動で満たされる。作成: `data/synth_train_sfens.jsonl` (15,000)、`data/synth_val_sfens.jsonl` (2,500)。 +2. **`scripts/data/subsample_synthetic_sfens.py`**: 15,000 は多いので **希少局面を残したまま 10,000 に間引き**。希少判定は「どこかの slot で count ≥ 3」OR「飛 / 角が hand に入っている」の OR 条件。実測では 14,906 件（99.4%）が rare 判定に入り、残り 94 件の「完全に静かな局面」だけがマジョリティから抜けた。rare 判定内は uniform で subsample。原本は `data/synth_train_sfens.jsonl.bak` に退避済み。 +3. **piyo-hook で 4 端末撮影**: `data/synth_train_sfens.jsonl` / `data/synth_val_sfens.jsonl` の全 SFEN を 4 端末でキャプチャ。1 SFEN × 4 端末 = 40,000 webp（train）+ 10,000 webp（val）。HF 既存分と合わせて `data/ocr//.webp` に配置。 +4. **`scripts/data/build_v3_manifest.py`**: HF cache の parquet から `sfen, hash, type` だけを pyarrow で吸い出し（画像列は触らない）、HF train 18k を同じ rare-preserving ルールで 10k に間引き、synth 10k と concat して `data/ocr_v3/train.jsonl` (20,000 行) を書き出す。val は subsample せず HF 2k + synth 2.5k = 4,500 行を全部残す。hash 重複は natural 側優先（実機 webp を持っている方を残す）。 + +### 修正結果（v3 train 20k、`data/ocr_v3/` の実測値） + +| 指標 | v2 (HF 18k) | **v3 train 20k (実測)** | v3 - v2 | +|---|---:|---:|---:| +| count=0 の比率 | 65.68% | **54.03%** | **-11.6pt** | +| count=1 の比率 | 21.20% | 24.25% | +3.0pt | +| count=3 以上の合計 | ~4% | **11.41%** | +7.4pt | +| count=7 の比率 | ~0% | 0.55% | +0.55pt | +| count=10 の比率 | 0.12% | 0.28% | +0.16pt | +| count=15 の比率 | ~0 | 0.07% | +0.07pt | +| count=18 の観測 | 4 件 | **17 件** | +13 | +| 飛の非ゼロ率 (S:R / G:R) | 17.8% / 17.1% | **30.1% / 28.5%** | 約 +12pt | +| 角の非ゼロ率 (S:B / G:B) | 42.9% / 41.5% | 44.5% / 43.6% | ほぼ同 | +| 歩の非ゼロ率 (S:P / G:P) | 68.0% / 64.0% | **78.4% / 77.1%** | +10〜13pt | + +「常に 0 ベースライン」の slot_acc が **65.7% → 54.0%**（-11.6pt）。pure synth 10k で得られる -20pt には届かないが、**サンプル数は 20k に増え、画像スタイルは実機で統一**。当初の見込み ~58.5% より 4.5pt 良かったのは、HF 18k の rare-preserving subsample がよく効いたため（18k 中 15,678 件（87%）が rare 判定に入り、common 側の「完全に静かな 8000 件」だけが落ちた）。 + +**v3 val の分布（4,500 行、`data/ocr_v3/val.jsonl` 実測）**： + +| 指標 | v2 val (HF 2k) | **v3 val 4.5k (実測)** | +|---|---:|---:| +| count=0 の比率 | 75.10% | **58.88%** | +| count=3 以上の合計 | 2.98% | 9.61% | +| count=11 以上の観測 | **0 件** | 通算 335 件（count=11 で 120、count=18 で 5） | +| count=18 の観測 | 0 件 | **5 件** | + +**val の count=11+ 問題が解消**され、regression head や高枚数向け class weight の効果を `val/hand/slot_acc` の per-count recall で直接評価できるようになった。 + +### 20k の内訳（provenance タグ、`type` field） + +`build_v3_manifest.py` の実行結果より： + +| type | 件数 | 由来 | +|---|---:|---| +| `synthetic` | 12,770 | 今回作った synth 10k + HF に元々含まれていた synth の残り | +| `existing` | 3,658 | HF の既存撮影済み `data/detector/*` reuse 分 | +| `natural` | 2,643 | HF の mate 系（`assets/mate{3,5,7,9,11}.sfen`）由来 | +| `opening` | 929 | HF の opening 系（`assets/start_sfens_ply{24,32}.txt`）由来 | +| **合計** | **20,000** | | + +per-source per-count recall を取れば、synth が hand tail の学習にどれだけ効いたかを直接測れる。 + +### v3 で調整するパラメータ + +**必要な変更（データマージに伴う）**： + +- **train データソース**: `--hf-repo-id` の HF loader だけでは synth 分が読めない。実装コストの低い順に： + - **(a) 新 HF split を publish**: `build_paired_to_hf.py` を synth 撮影後の webp も拾うよう拡張し、`ocr_paired_v3` として push。既存の HF loader をそのまま使える。**推奨**。`data/ocr_v3/train.jsonl` と `data/ocr_v3/val.jsonl` はこの入口に渡す形。 + - **(b) train_board_ocr.py に extra jsonl の入口**: `--extra-train-jsonl` / `--extra-val-jsonl` を追加、HF pool と concat する `ConcatDataset` にする。HF publish 手間を省ける代わりに loader 側に merge ロジックが増える。 +- **`--class-weight-clip-max 10.0 → 15.0`**: v3 20k 実測の raw class weight は count=10 で 19.1、count=15 で 78.4、count=18 で 867。旧 clip=10 だと **count≥7 が全部同じ重み**でキャップされる。**clip を 15 に緩める** と sqrt 後 count=7 で 3.10、count=10 で 4.37、count=13 で 6.62、count=15 で 8.85、count=17 以上が 15 でキャップ、と個別に立ち上がる。それより上は継続キャップ。 +- **`EPOCHS=60 → 50`**: サンプル数が 18k → 20k で 11% 増、batch=128 では **epoch あたり 141 → 156 ステップ**、v2 の 60 epoch (8460 steps) 相当は約 **54 epoch** で並ぶ。**50 epoch（v2 比 -8% 総ステップ、cosine scheduler で後半の実効 lr を下げるぶんはトントン）** に設定して cosine と組で回す。 + +**推奨追加（既存の TODO を v3 で入れると綺麗）**： + +- **§D LR scheduler の導入**: `AdamW(lr=6e-4)` 固定を `CosineAnnealingLR(T_max=EPOCHS)` に。resume 時は scheduler state も ckpt に含める。データが変わって baseline を再取得する v3 は、scheduler も同時投入して次の baseline に組み込む好機。 +- **`--hand-mode regression` の CLI 露出**: model 側の実装 (`board_ocr.py`) は既にあるが `train_board_ocr.py` に CLI 引数が無い。マージ後の分布は count=3〜18 が実在するので、regression head の効果測定は v2 データより格段にやりやすい。**別実験として v3-cls / v3-reg の 2 系統を並行して回すのが理想**。 + +**据置き**： + +- `BATCH_SIZE=128` / `IMAGE_SIZE=288` / `BACKBONE=mobilenet_v3_small`: v2 との direct delta を測るため、まずはこの 3 点を据え置き、データ差 + class weight + epoch + scheduler の効果を分離する。backbone / image_size の bump は v3 の結果を見てから、v4 で入れる。 +- **`hand_weight=1.0`**: v2 で入れた設定、そのまま。 +- **X-1 hand logit マスキング**: そのまま。 + +### v3 想定 config + +```bash +# 前提: HF に ocr_paired_v3 が publish 済み、あるいは train_board_ocr.py に extra jsonl 入口がある。 +EPOCHS=50 \ +BATCH_SIZE=128 \ +LR=6e-4 \ +IMAGE_SIZE=288 \ +BACKBONE=mobilenet_v3_small \ +CKPT_DIR=./runs/board-ocr-v3 \ +HF_REPO_ID=ultemica/piyoshogi \ # または v3 split の repo +./scripts/train.sh \ + --class-weight-clip-max 15.0 \ + --lr-scheduler cosine # 実装後 +``` + +**目標**：sfen_full > 0.60（v2 baseline を追加 hand 分布で越える）。同時に per-count recall（特に count=5〜18）を `scripts/inspect/diagnose_hand.py` で計測して、v2 との hand 側改善差を数字で押さえる。**per-source (natural / synthetic) の per-count recall** も並行して見ると、synth の効きを直接評価できる。 + +### 開始前チェックリスト + +1. **piyo-hook 撮影完了待ち**: `data/synth_train_sfens.jsonl` + `data/synth_val_sfens.jsonl` の全 SFEN が 4 端末で撮影されて `data/ocr//.webp` に落ちること。所要時間はデータ量次第（10k+2.5k SFEN × 4 端末 = 50,000 webp）。 +2. **manifest 統合（済）**: `scripts/data/build_v3_manifest.py` を実行済み。`data/ocr_v3/train.jsonl` (20,000 行) と `data/ocr_v3/val.jsonl` (4,500 行) が生成済み。この 2 本を parquet 化担当に渡して `ocr_paired_v3` に育ててもらう。 +3. **HF publish or extra jsonl 対応**: (a) `build_paired_to_hf.py` を `data/ocr_v3/` 入力に切り替えて回し `ocr_paired_v3` として push、または (b) `train_board_ocr.py` に `--extra-train-jsonl` の入口を追加。 +4. §D CosineAnnealingLR の実装 + resume 対応。 +5. `--hand-mode` CLI 引数の追加（別実験線として）。 +6. `--class-weight-clip-max` を 15 に上げても勾配が暴発しないことを 5 epoch で確認。 +7. baseline 比較のため v2 の epoch 60 ckpt を凍結（`runs/board-ocr-v2/` を触らない）。 diff --git a/docs/README.md b/docs/README.md index d7d0996..2d3acc7 100644 --- a/docs/README.md +++ b/docs/README.md @@ -10,6 +10,8 @@ MITO ↔ mito-train 契約と mito-train 固有ドキュメント。 ## mito-train 固有 - `piyo-piece-templates.md` — 駒テンプレの命名規則 (14 デザイン × 30 駒) +- `backbones.md` — BoardOCR がサポートする 9 バックボーンの params / 用途 / 配信先まとめ +- `ocr-scaling-outlook.md` — sweep 結果からのエポック数・画像サイズ増の効き見込み ## データセットの正規 diff --git a/docs/backbones.md b/docs/backbones.md new file mode 100644 index 0000000..ecab32a --- /dev/null +++ b/docs/backbones.md @@ -0,0 +1,368 @@ +# Backbones + +BoardOCR がサポートする 9 個のバックボーンを、パラメータ数・想定用途・配信先ごとに整理する。全モデル `BoardOCR` から `--backbone ` で切替可能。 + +## 一覧 + +| Backbone | Params | fp32 | fp16 | int8 | 出典 | 特徴 | +|---|---|---|---|---|---|---| +| [mobilenet_v3_small](#mobilenet_v3_small) | 1.1M | 4.4 MB | 2.2 MB | ~1.1 MB | torchvision | 最軽量 | +| [mobilenet_v3_large](#mobilenet_v3_large) | 3.3M | 13 MB | 6.5 MB | ~3.3 MB | torchvision | small の上位 | +| [convnext_atto](#convnext_atto) | 3.5M | 14 MB | 7 MB | ~3.5 MB | timm | ConvNeXt最小 | +| [efficientnet_b0](#efficientnet_b0) | 4.4M | 18 MB | 9 MB | ~4.4 MB | torchvision | 定番中庸 | +| [convnext_femto](#convnext_femto) | 4.9M | 20 MB | 10 MB | ~4.9 MB | timm | ConvNeXt 次点 | +| [efficientnet_b1](#efficientnet_b1) | 6.9M | 28 MB | 14 MB | ~6.9 MB | torchvision | b0 の上位 | +| [convnext_pico](#convnext_pico) | 8.7M | 35 MB | 17 MB | ~8.7 MB | timm | ConvNeXt 小 | +| [convnext_nano](#convnext_nano) | 15.1M | 61 MB | 30 MB | ~15.1 MB | timm | Pareto中庸 | +| [convnext_tiny](#convnext_tiny) | 28.0M | 112 MB | 56 MB | ~28.0 MB | torchvision | モバイル上限 | + +数値は `pretrained=False` で `BoardOCR(backbone=...)` を組んだときのパラメータ総数。int8 は「近似バイト数」であり、実際は量子化スキームで前後する。 + +## 用途マトリクス + +| 配信先 | 推奨 backbone | 理由 | +|---|---|---| +| Cloudflare Workers + workers-wonnx (WebGPU) | mobilenet_v3_small〜convnext_atto | R2 にモデル置いて WebGPU 経由。CPU 推論より現実的 | +| Cloudflare Workers 純 CPU 推論 | mobilenet_v3_small (int8) | GPU 無しで実用ラインギリ、Paid tier 必須 | +| Cloudflare Workers AI（ホスト済み） | 該当なし（現状） | カタログ限定、BYOM は 2026 半ば時点で未 GA | +| ブラウザ（初回DL重視） | mobilenet_v3_small / large, convnext_atto | int8 で数MB、初回ロード軽い | +| ブラウザ（精度重視） | convnext_nano | int8 15MB、精度/サイズのバランス最良 | +| 2024+ フラグシップスマホ | convnext_tiny | NPU性能余裕、精度上限狙える | +| 2024+ ミッドレンジスマホ | convnext_nano | 旧機種互換も考慮した実用ライン | +| ネイティブアプリ（iOS/Android） | convnext_tiny | 精度上限を狙える現実的サイズ | +| サーバ推論のみ | convnext_tiny | 拡張余地あり（convnext_small まで検討可） | +| 高速オンデバイス推論 | efficientnet_b0 | Compound scaling で少パラでも精度良 | +| 実験のベースライン | mobilenet_v3_small | 学習速度・容量ともに軽く鉄板 | + +## モデル別詳細 + +### mobilenet_v3_small + +- **Params**: 1.1M / **out channels**: 576 +- **出典**: `torchvision.models.mobilenet_v3_small` +- **アーキ**: Inverted residuals + Squeeze-and-Excite + h-swish activation +- **特徴**: 2019年 Google 設計。Neural Architecture Search で mobile 向けに直接最適化。最軽量クラス。 +- **強み**: 学習/推論ともに超高速、int8 で 1MB切る、ブラウザで一瞬でロード。 +- **弱み**: 特徴容量が小さく、細部識別（成香/成桂/と金 の漢字区別）が伸びにくい。 +- **想定用途**: パイプライン検証、smoke test、ブラウザで「とりあえず動く」ライン。 + +### mobilenet_v3_large + +- **Params**: 3.3M / **out channels**: 960 +- **出典**: `torchvision.models.mobilenet_v3_large` +- **アーキ**: small と同系統、より広いチャネル・深い block 構成 +- **特徴**: mobilenet_v3_small のスケールアップ版。同じ設計思想で 3 倍の容量。 +- **強み**: small の学習パイプラインそのまま流用可、int8 で 3MB とまだ軽い。 +- **弱み**: efficientnet_b0 と比べると精度/パラの効率で若干劣る（新しい設計思想の差）。 +- **想定用途**: small の精度で足りない時の**最小限のスケールアップ**。 + +### convnext_atto + +- **Params**: 3.5M +- **出典**: `timm.create_model("convnext_atto")` +- **アーキ**: 4-stage ConvNeXt を極限まで縮小（depth=[2,2,6,2], dims=[40,80,160,320]） +- **特徴**: 2022年 Facebook AI の ConvNeXt を極小化。LayerNorm + GELU + depthwise 7x7 conv の現代アーキ。 +- **強み**: ConvNeXt 系の効率良い設計を最軽量帯に持ってきたもの。BN 不使用なので DDP でも SyncBN 気にせず。 +- **弱み**: timm 依存。 +- **想定用途**: mobilenet_v3_large と同じサイズ帯で ConvNeXt アーキを試したい時。 + +### efficientnet_b0 + +- **Params**: 4.4M / **out channels**: 1280 +- **出典**: `torchvision.models.efficientnet_b0` +- **アーキ**: Compound scaling で width/depth/resolution を同時に最適化した MBConv ベース +- **特徴**: 2019年 Google、当時 SOTA の精度/パラ効率。後継の EfficientNet-V2 もあるが b0 の完成度が高く定番として残る。 +- **強み**: 4.4M で ImageNet top-1 77% 級。この容量帯のリファレンス実装。 +- **弱み**: SE ブロックの計算がやや重い。DWConv の並列度が若干控えめ。 +- **想定用途**: 「小さいけど精度も欲しい」の第一候補。ブラウザ配信に強い。 + +### convnext_femto + +- **Params**: 4.9M +- **出典**: `timm.create_model("convnext_femto")` +- **アーキ**: ConvNeXt をやや大きくした版（depth=[2,2,6,2], dims=[48,96,192,384]） +- **特徴**: atto と pico の間のサイズ。ConvNeXt 系の精度/パラのスイートスポット候補。 +- **強み**: LayerNorm ベースで DDP フレンドリー。int8 で 5MB。 +- **想定用途**: efficientnet_b0 のカウンタパート実験。ConvNeXt vs EfficientNet の比較の重要ピボット。 + +### efficientnet_b1 + +- **Params**: 6.9M +- **出典**: `torchvision.models.efficientnet_b1` +- **アーキ**: b0 を compound scaling で 1 段引き上げた版 +- **特徴**: 入力解像度 240x240 で pretrained（BoardOCR では 224 リサイズで使用）。 +- **強み**: b0 より数%高い精度、それでも 7M で mobile 圏。 +- **弱み**: b0 との差が僅かで、コスト比の効率悪化傾向。 +- **想定用途**: b0 で足りず nano まで行きたくない時のギャップ埋め。 + +### convnext_pico + +- **Params**: 8.7M +- **出典**: `timm.create_model("convnext_pico")` +- **アーキ**: ConvNeXt（depth=[2,2,6,2], dims=[64,128,256,512]） +- **特徴**: 10M 前後の中量級 ConvNeXt。 +- **強み**: femto の 2 倍の容量、精度がぐっと上がる帯。 +- **想定用途**: mobile と server の中間帯を狙うとき。ブラウザだと少し重いが実用範囲。 + +### convnext_nano + +- **Params**: 15.1M +- **出典**: `timm.create_model("convnext_nano")` +- **アーキ**: ConvNeXt（depth=[2,2,8,2], dims=[80,160,320,640]） +- **特徴**: ConvNeXt-Tiny を「肉薄しつつ半分のサイズ」に絞ったもの。 +- **強み**: **Pareto Front の中央**。int8 で 15MB、convnext_tiny の精度に大きく譲らない可能性。 +- **弱み**: timm 依存。 +- **想定用途**: **ブラウザ配信の本命候補**。精度が求められて 30MB の初回DLが許容できるならこれ。 + +### convnext_tiny + +- **Params**: 28.0M / **out channels**: 768 +- **出典**: `torchvision.models.convnext_tiny` +- **アーキ**: 4-stage ConvNeXt（depth=[3,3,9,3], dims=[96,192,384,768]） +- **特徴**: 2022年 Facebook AI の ConvNeXt-Tiny。Transformer 系に迫る精度を ConvNet で達成。 +- **強み**: 精度上限。ImageNet top-1 82%。torchvision 標準実装で ecosystem 対応◎。 +- **弱み**: int8 でも 28MB でモバイル配信の実用上限、旧機種でメモリ厳しめ。fp16 だと 56MB。 +- **想定用途**: **モバイル配信の上限**。ネイティブアプリなら現実解。ブラウザは fp16/int8 前提。 +- **注意**: これより大きい convnext_small (50M) はブラウザには重い。ネイティブでも flagship 機限定。 + +## デプロイ先別の詳細 + +### Cloudflare Workers（自前ONNXモデル推論） + +2026年時点、Cloudflareで自前 ONNX を動かす選択肢は3つ。それぞれ制約が違うので使い分けが必要〜。 + +#### ルート A: `workers-wonnx` パターン（推奨） + +公式サンプル [`cloudflare/workers-wonnx`](https://github.com/cloudflare/workers-wonnx) が示す方法。WebGPU ベースの ONNX ランタイム WONNX を Worker 上で動かし、モデル本体は R2 バケットから fetch する。 + +- **GPU 使えるので推論はそこそこ速い**（WebGPU 経由） +- **モデルサイズ制約回避**: R2 に置くから Worker のバンドルサイズ制限に縛られない +- **BoardOCR の場合の実用度**: convnext_tiny クラスは重すぎるので、**mobilenet_v3_small〜convnext_atto が現実的** +- 推論時間目安: mobilenet_v3_small int8 で **20〜50ms**、convnext_atto で **80〜200ms**（WebGPU 経由） + +#### ルート B: Workers AI カタログ（該当モデル無し） + +Cloudflare 側で GPU 実行してくれる仕組み。ただし利用可能モデルはカタログ限定で、**カスタム ONNX は現状 GA していない**（BYOM がロードマップにはいるがまだ未提供、2026年半ば時点）。 + +- カタログにあるのは Llama, Whisper, Stable Diffusion 系 +- **将棋盤 OCR 用途は該当なし** → 使えない +- 将来 BYOM 対応したら convnext_tiny クラスまで載る想定 + +過去に「Constellation」というカスタム ONNX 実行サービスがあったが、Workers AI に統合され独立プロダクトとしては終息。 + +#### ルート C: Worker CPU 上で ONNX Runtime Web / WASM + +昔ながらの方法。**GPU 使えないので実用性はかなり厳しい**。 + +- **CPU 時間**: Free 10ms（推論不可）/ Paid tier で 30秒 (bundled) / 5分 (unbound) +- **メモリ**: 128MB +- **推論時間の見積もり**（Worker CPU で 224x224 1枚推論） + - mobilenet_v3_small int8: **200〜400ms** — Paid tier ならぎりぎり + - efficientnet_b0 int8: **500ms〜1s** — 苦しい + - convnext_atto 以上: **1s超** — Unbound のみ実用 + +### Workers 料金の目安（2026年時点） + +- **Workers Paid**: 月$5、10M req + 30M CPU-ms 込み +- 超過分: 1M req あたり$0.30、1M CPU-ms あたり$0.02 +- **Workers AI**: 1,000 Neurons あたり$0.011（無料枠1日10,000 Neurons） + - Neurons = リクエスト実行に必要な GPU コンピュートの単位 + - カスタム ONNX（`workers-wonnx`）は Workers AI ではなく Workers 側の課金体系 +- **R2 ストレージ**: モデルファイル置き場、10GB/月まで無料、以降 $0.015/GB/月 + +**BoardOCR のケーススタディ** +- mobilenet_v3_small int8 (1.1MB) を R2 に置いて workers-wonnx で提供 +- 1回の推論: ~30ms CPU + R2 fetch（キャッシュ効くから初回のみ） +- 1M リクエスト/月 想定: Worker Paid $5 + CPU 超過分ほぼゼロ = **~$5/月で運用可能** +- ただし GPU コンピュート単価は今後変動可能性あり、公式ドキュメント要確認 + +### 2024年以降のスマートフォン + +処理性能が急速に上がり、モデル選択の自由度が上がった帯。 + +#### iPhone 15 Pro / iPhone 16（A17/A18/A18 Pro） + +- **Neural Engine**: 17〜38 TOPS +- **RAM**: 8GB +- **推論経路**: CoreML → Neural Engine or Metal GPU +- **推論時間**（BoardOCR 224x224 fp16 想定） + - convnext_tiny: **~15〜30ms** + - convnext_nano: **~10〜20ms** + - mobilenet_v3_small: **~3〜8ms** +- **上限モデル**: convnext_tiny 余裕、**convnext_small (50M) も動く** +- **推奨**: convnext_tiny（精度優先） or convnext_nano（アプリ容量抑えたい） + +#### iPhone 15 / 16（A16/A18 non-Pro） + +- **Neural Engine**: 15.8〜17 TOPS +- **RAM**: 6〜8GB +- Pro とほぼ同等の推論時間 +- 上限は convnext_tiny、convnext_small はメモリ的にギリ + +#### Snapdragon 8 Gen 3 / 8 Gen 4（2024年 Android フラグシップ） + +- Galaxy S24, Xiaomi 14 系列 +- **NPU (Hexagon)**: 45+ TOPS +- **推論経路**: NNAPI / QNN / ONNX Runtime Mobile + XNNPACK +- **推論時間**（BoardOCR 224x224 fp16 想定） + - convnext_tiny: **~20〜40ms** + - convnext_nano: **~15〜25ms** +- iPhone 15 Pro クラスと同程度の能力 +- 推奨: convnext_tiny + +#### Google Pixel 9 / 9 Pro（Tensor G4） + +- **TPU (Edge TPU 派生)**: 高い量子化推論性能 +- int8 量子化と相性◎ +- convnext_tiny int8: **~30〜50ms** +- 推奨: convnext_tiny (int8) + +#### 2024年ミッドレンジ Android（Snapdragon 7 Gen 3 / Dimensity 7300 級） + +- **NPU**: 10〜20 TOPS +- **推論時間** + - convnext_tiny: **~80〜150ms**（許容範囲だが体感重くなる） + - convnext_nano: **~40〜80ms** + - efficientnet_b0: **~20〜40ms** +- **推奨**: convnext_nano。convnext_tiny は動くが体感速度が犠牲になるケースあり + +#### まとめ表（2024+ スマホ想定） + +| 端末クラス | 推奨 backbone | 推論時間 | 備考 | +|---|---|---|---| +| iPhone 15/16 Pro | convnext_tiny | ~20ms | 精度優先で余裕 | +| iPhone 15/16 (無印) | convnext_tiny | ~30ms | 標準的選択 | +| Snapdragon 8 Gen 3+ | convnext_tiny | ~30ms | Android旗艦 | +| Pixel 9 (Tensor G4) | convnext_tiny (int8) | ~40ms | int8 との相性◎ | +| ミッドレンジ2024 | convnext_nano | ~40〜80ms | tiny は体感重い | +| 旧機種互換重視 | convnext_atto / mobilenet_v3_large | ~30〜100ms | 2020年頃のミッド機種でも動く | + +**共通の実装ポイント** +- **CoreML / ONNX 変換**: `torch.onnx.export` → 端末側 SDK でロード +- **fp16 量子化**: モバイル NPU/GPU で最も相性が良い、精度低下は無視できる +- **int8 量子化**: サイズと速度に効くが、キャリブレーションデータ必要 +- **画像入力の pre-processing**: `LongestMaxSize + PadIfNeeded + Normalize` を端末側で再現。ImageNet mean/std をハードコード + +## 推論時間の比較表（横断） + +224x224 入力・1画像あたりの推論時間の**目安**。数値は他モデルの公開ベンチと BoardOCR のパラメータ規模から推定した見積もりで、実測値ではない。 + +| Backbone | Worker CPU 推論 | Worker + workers-wonnx (WebGPU) | iPhone 15/16 Pro (NE) | Snapdragon 8 Gen 3 (NPU) | Pixel 9 (TPU int8) | ミッドレンジ Android 2024 | +|---|---|---|---|---|---|---| +| mobilenet_v3_small | **200〜400ms** | 20〜50ms | 3〜8ms | 5〜10ms | 5〜10ms | 15〜30ms | +| mobilenet_v3_large | 400〜800ms | 40〜100ms | 5〜12ms | 8〜15ms | 8〜15ms | 25〜50ms | +| convnext_atto | 800ms〜2s | 80〜200ms | 6〜12ms | 10〜18ms | 10〜18ms | 30〜60ms | +| efficientnet_b0 | 500ms〜1s | 50〜120ms | 5〜10ms | 8〜15ms | 8〜15ms | 20〜40ms | +| convnext_femto | 1〜3s | 100〜250ms | 8〜15ms | 12〜20ms | 12〜20ms | 40〜80ms | +| efficientnet_b1 | 800ms〜1.5s | 80〜180ms | 8〜15ms | 12〜20ms | 12〜20ms | 30〜60ms | +| convnext_pico | 2〜5s | 200〜500ms | 10〜18ms | 15〜25ms | 15〜25ms | 60〜120ms | +| convnext_nano | 5〜10s | 400ms〜1s | 10〜20ms | 15〜25ms | 15〜30ms | 40〜80ms | +| convnext_tiny | 10〜30s | 800ms〜2s | 15〜30ms | 20〜40ms | 30〜50ms | 80〜150ms | + +**読み取りポイント** +- **Worker CPU 推論**: convnext_atto より重いモデルはリクエスト毎の CPU 時間制約でほぼ実用不可 +- **workers-wonnx (WebGPU)**: CPU 推論の 10〜20倍速。convnext_atto までは実用範囲 +- **モバイル NPU/TPU 経由**: convnext_tiny でも 30〜50ms、体感 60fps は無理でも 15〜20fps は出る +- **ミッドレンジスマホでのラインは convnext_nano**: tiny だと 100ms 超で操作もっさり +- サーバ GPU 推論なら全 backbone が数ms〜十数ms で完結（比較対象外だが参考として） + +## 精度と実用性の関係 + +「val sfen_full が何%なら実用に耐えるか」は用途で決まる。以下は将棋盤 OCR で想定される要件レベル。 + +### 用途別の実用ライン + +| 用途 | 必要 sfen_full acc | 理由 | +|---|---|---| +| **カジュアル閲覧**（画像から盤面表示、参考程度） | **~85%以上** | 誤認識してもユーザーが手動確認できる | +| **プレイ再現・棋譜記録** | **~95%以上** | 手が進むごとに再入力面倒、8割成功でも運用崩壊 | +| **エンジン解析への入力** | **~99%以上** | 盤面1マス違えば局面全く変わる。プロ棋士は絶対に許容しない | +| **公式棋譜化・研究用途** | **~99.5%以上** | 人手校正コスト削減の意味を成すライン | + +### cell_acc と sfen_full の関係（重要） + +`sfen_full` は「盤面81マス全部 + 持ち駒14スロット全部」正解の率。**セル単位の精度がわずかに下がるだけで sfen_full は急落する**。 + +理論値（独立仮定）： + +| board cell_acc | 理論 sfen_full (81マス独立仮定) | +|---|---| +| 95% | 1.5% | +| 97% | 8.5% | +| 99% | 44% | +| 99.5% | **66%** | +| 99.8% | **85%** | +| 99.9% | **92%** | + +**実測はこの理論値より高くなる**（同じ画像内での間違いが空間相関するため）。ただし cell_acc 97% と 99% では sfen_full が 数十%〜数倍差になる。 + +**教訓**: cell_acc を 99%台後半まで押し上げないと、実用ラインに乗らない。 + +### 何が精度を押し上げるか + +sfen_full を伸ばす手段の効き順（BoardOCR での経験則）： + +1. **image_size の増（224 → 288 → 384）**: 一番効きやすい。9x9 マス識別に解像度が直接効く +2. **モデル容量の増（backbone を上げる）**: 特に細部識別（成香 vs 成桂、と金 vs 金）で効く +3. **augmentation の適正化**: SNS 劣化への頑健性（cell_acc への直接寄与は小さいが実運用差で効く） +4. **hand_mode = "regression"** への切替: 持ち駒枚数のオーディナル情報を活かす +5. **エポック数の増**: 逓減あるが、cell_acc 97% → 98% への最後の詰めで効く +6. **蒸留**: 上限に近付いた後の最後の押し上げ手段 + +### 実用的な妥協点 + +| 想定シナリオ | 現実的な組み合わせ | +|---|---| +| 「まず動くもの」 | mobilenet_v3_small, image_size=224, sfen_full ~50%狙い | +| 「棋譜記録に使える」 | convnext_nano, image_size=288, sfen_full ~90%狙い | +| 「エンジン解析まで」 | convnext_tiny, image_size=384, sfen_full ~99%狙い | +| 「公式棋譜化」 | convnext_tiny + 蒸留, image_size=384, sfen_full 99.5%+ | + +### 実運用での妥協策 + +100% は理論上不可能に近い。運用側で吸収する仕組みも視野に： + +- **信頼度スコア表示**: 「このマスは確信度低め」と UI で示し、人手校正を促す +- **候補提示**: argmax だけでなく top-3 を出す、間違い時に選び直しやすく +- **文脈補正**: 「将棋のルール上、この位置に相手の玉は無い」等の後処理 +- **A/B の複数モデル投票**: convnext_tiny + efficientnet_b0 の合議で頑健性↑ + +## 選び方の判断樹 + +``` +用途は？ +├─ 実験/検証 → mobilenet_v3_small +├─ ブラウザ配信 +│ ├─ 精度優先 → convnext_nano (int8 15MB) +│ └─ サイズ優先 → mobilenet_v3_large or efficientnet_b0 (int8 3-5MB) +├─ ネイティブアプリ → convnext_tiny (int8 28MB) +└─ サーバ推論のみ + ├─ 十分なら → convnext_tiny + └─ さらに精度 → 蒸留 or image_size 増 (backboneはtiny固定) +``` + +## パラメータ以外に効く要素 + +同じ backbone でも以下で精度が変わる： + +- **image_size**: 224 → 384 で cell_acc 数%改善する可能性。将棋盤 9x9 の識別は解像度律速。 +- **hand_mode**: `classification`（デフォ）vs `regression`。持ち駒枚数のオーディナル情報を活かすなら regression。 +- **augmentation**: SNS 圧縮/ノイズを想定して調整可（`capture_dataset.py` の `build_transform`）。 +- **preload**: 学習速度に効くが精度には無影響。 + +## 蒸留 (Knowledge Distillation) の余地 + +「convnext_tiny では足りないがサイズは上げたくない」場合の常道： + +1. サーバで convnext_small / base を teacher として学習 +2. mobile 向け convnext_nano / tiny を student として、teacher の soft label で学習 +3. パラメータ数を保ったまま数%精度向上 + +BoardOCR では未実装だが、モバイル配信で精度を詰めたい場合の次のカードとして有力。 + +## 実測データ（今後追加） + +各 backbone の val cell_acc / sfen_full_acc / wall-clock は sweep 完了後に W&B で並び、ここに引用予定。 + +- W&B project: `mito-train-board-ocr` +- run 名は各 backbone 名そのまま（例: `mobilenet_v3_small`, `convnext_tiny`） diff --git a/docs/ocr-scaling-outlook.md b/docs/ocr-scaling-outlook.md new file mode 100644 index 0000000..a4fc19f --- /dev/null +++ b/docs/ocr-scaling-outlook.md @@ -0,0 +1,272 @@ +# OCR スケーリング見通し（エポック数・画像サイズを増やしたら？） + +feat/ddp-and-backbone-sweep で回した 9 バックボーンの sweep 結果（2026-07-12）を材料に、エポックを伸ばした場合・入力解像度を上げた場合に精度がどこまで伸びそうかを見積もる。 + +前提の指標定義は [`ocr-metrics.md`](./ocr-metrics.md)、モデル選択の全体像は [`backbones.md`](./backbones.md) を参照。ここでは sweep から出た **数字だけ** を根拠に、次のイテレーションの方針を出すのがゴール。 + +## sweep の共通条件 + +`scripts/train_backbones.sh` のデフォルト（2026-07-12 現行）： + +| 項目 | 値 | +|---|---| +| epochs | 50 | +| batch size | 32 | +| image size | **224** | +| lr | 3e-4 (AdamW 固定、scheduler 無し) | +| num workers | 16, prefetch 8, `--preload` | +| dataset | `ultemica/piyoshogi` (mode=full) | +| hand head | logit マスキング済み (v2)、hand_weight=1.0、class weight sqrt+clip | +| val 測定間隔 | 2 epoch ごと | + +同一条件下で backbone だけを差し替えた ちょい厳密な比較になっている。lr scheduler も EMA も distillation も入っていない **素の baseline**。 + +## 各バックボーンの val 数値 (epoch 50 時点) + +| Backbone | Params | val cell_acc | val sfen_full | train sfen_full | train loss | +|---|---:|---:|---:|---:|---:| +| mobilenet_v3_small | 1.1M | 0.966 | 0.385 | 0.185 | 0.4549 | +| mobilenet_v3_large | 3.3M | 0.989 | 0.593 | 0.428 | 0.1288 | +| convnext_atto | 3.5M | 0.991 | 0.611 | 0.541 | 0.0828 | +| efficientnet_b0 | 4.4M | 0.991 | 0.661 | 0.483 | 0.1128 | +| convnext_femto | 4.9M | 0.992 | 0.645 | 0.589 | 0.0635 | +| efficientnet_b1 | 6.9M | 0.992 | 0.663 | 0.475 | 0.1093 | +| convnext_pico | 8.7M | 0.992 | 0.634 | 0.632 | 0.0520 | +| convnext_nano | 15.1M | 0.993 | 0.660 | 0.680 | 0.0436 | +| **convnext_tiny** | **28.0M** | **0.994** | **0.716** | 0.659 | 0.0560 | + +- `val cell_acc` は **どのモデルも 0.966〜0.994** で頭打ち感。差は 3pt 弱。 +- `val sfen_full` は **0.385〜0.716** で **33pt** も開く。プライマリ指標はまだまだ動く。 +- convnext_nano は再学習して 50 epoch 完走（run `nyrsx25c`）。0.660 で **efficientnet_b0 (0.661) とほぼ同点**、pico (0.634) からは +2.6pt、tiny (0.716) には -5.6pt 届かず。15.1M の割に伸び切らず、Pareto 上では efficientnet_b0 (4.4M) と重なる位置。**train sfen 0.680 > val 0.660** で train 側の伸び余地は残っている。 + +train と val のギャップは全モデルで **val > train**（board head は val が優勢、hand head は val の方が難しい局面が来る）。過学習の兆候は現時点ゼロで、まだエポックを伸ばしてよい状態。 + +## val sfen_full の epoch 別トラジェクトリ + +val は 2 epoch おき測定なので、代表点として ep10 / 20 / 30 / 40 / 50 を並べる。 + +| Backbone | ep10 | ep20 | ep30 | ep40 | ep50 | ep40→50 差分 | +|---|---:|---:|---:|---:|---:|---:| +| mobilenet_v3_small | 0.241 | 0.334 | 0.367 | 0.388 | 0.385 | **-0.003** | +| mobilenet_v3_large | 0.389 | 0.491 | 0.537 | 0.568 | 0.593 | +0.025 | +| convnext_atto | 0.482 | 0.562 | 0.599 | 0.611 | 0.611 | 0.000 | +| efficientnet_b0 | 0.447 | 0.570 | 0.611 | 0.640 | 0.661 | **+0.021** | +| convnext_femto | 0.529 | 0.577 | 0.612 | 0.643 | 0.645 | +0.002 | +| efficientnet_b1 | 0.428 | 0.526 | 0.596 | 0.637 | 0.663 | **+0.026** | +| convnext_pico | 0.512 | 0.598 | 0.618 | 0.641 | 0.634 | -0.007 | +| convnext_nano | 0.586 | 0.626 | 0.640 | 0.647 | 0.660 | +0.013 | +| convnext_tiny | 0.620 | 0.690 | 0.710 | 0.699 | 0.716 | +0.017 | + +読み取れること： + +- **cell_acc は epoch 6〜10 で 0.98 台に到達、以後の伸びしろは 1pt 未満**。ボトルネックは常に hand head 経由の sfen_full。 +- **mobilenet_v3_small は 40 epoch で頭打ち**。1.1M パラの容量律速でここから伸ばしても線形改善は期待できない。 +- **convnext_atto と convnext_pico はサイズ違いなのに ep50 で 0.611 / 0.634 でほぼ団子**。ep40 でも同傾向。ConvNeXt 系はこの容量帯（3〜9M）で頭打ちしている可能性大。 +- **EfficientNet 系（b0, b1）は 40→50 で +2pt 以上伸びており、直近 10 epoch でも上げ足を残している**。 +- **convnext_nano は ep30→50 で +0.020、ep40→50 で +0.013 の緩やかな上昇**で飽和には至っていない。ただし ep50 で 0.660 は efficientnet_b0 (0.661) と同点で、pico からのスケーリング効率は悪い。ConvNeXt の 3.5M/4.9M/8.7M/15.1M で 0.611/0.645/0.634/0.660 と非単調、この容量帯全体が hand head 側で律速されている疑いが強い。 +- **convnext_tiny は ep30 で 0.710 に到達し ep40 でノイズで下がったが ep50 で再度 0.716**。5 epoch 移動平均を取ればまだ緩やかに上向き。飽和はまだ。 + +## エポックを伸ばしたときの見込み + +各モデルの ep30→50 の 20 epoch 分の増分から、**ep50→100 の追加 50 epoch でどれくらい伸びるか** を粗く外挿する。 + +前提： +- 現状は AdamW `lr=3e-4` **固定**（scheduler 無し）。後半になるほど lr が高すぎて振動する典型パターン。 +- v1 baseline（TRAINING_PLAN.md）でも同種の後半停滞が観察されている（v1 40 ep で sfen 0.486 → v2 hand 改善で ep50 0.385、条件が違う）。 +- ep30→50 で伸びた幅の **1/2〜1/3** が ep50→100 の追加改善の目安（logistic 型に減衰する経験則）。ここに lr scheduler を入れると +2〜3pt が乗る、というのが `TRAINING_PLAN.md#D` の期待値。 + +| Backbone | ep30→50 増分 | ep100 予想 (据置lr) | +cosine LR で狙える上振れ | 備考 | +|---|---:|---:|---:|---| +| mobilenet_v3_small | +0.018 | ~0.395 | ~0.42 | 容量律速。飽和目前 | +| mobilenet_v3_large | +0.056 | ~0.62 | ~0.65 | まだ上げ足あり | +| convnext_atto | +0.012 | ~0.62 | ~0.65 | 30 epoch で飽和気配 | +| efficientnet_b0 | +0.050 | ~0.69 | ~0.72 | 直近も上向き、伸びる | +| convnext_femto | +0.033 | ~0.66 | ~0.69 | 中庸 | +| efficientnet_b1 | +0.067 | **~0.70** | **~0.73** | 直近の勾配が最良 | +| convnext_pico | +0.016 | ~0.65 | ~0.67 | atto と同水準に張り付き | +| convnext_nano | +0.020 | ~0.68 | ~0.71 | 15.1M の割に伸び切らず、b0 と同格 | +| convnext_tiny | +0.006 | ~0.73 | **~0.76** | ばらつきあり、平均で微増 | + +これは **「同じ lr のまま 100 epoch まで回した場合」** の見込みで、実運用の伸び余地としては **cosine annealing / ReduceLROnPlateau を入れたときの +0.02〜0.03** をそこに乗せた側が現実的な上限。 + +### エポック増でも越えられない壁 + +sfen_full が上記予想値で頭打ちする理由： + +1. **cell_acc が既に 0.99 台に張り付いている**。ここから 1pt 押し上げても sfen_full にはあまり効かない（[`ocr-metrics.md` の Q1 表](./ocr-metrics.md)、および [`backbones.md` の cell_acc→sfen_full 表](./backbones.md#cell_acc-と-sfen_full-の関係重要)）。 +2. **hand_full_acc がまだ 0.7 前後**。sfen_full ≈ board_correct × hand_correct なので、hand を伸ばさない限り sfen は上に抜けない。詳細は下記「hand の失敗パターン」を参照。 + +つまり **エポックを 2 倍に伸ばしても sfen_full の到達値は +3〜5pt**。実用ライン 90% には遠く、エポック単体の追加投資では届かない。 + +## hand の失敗パターン + +hand head は構造上 **駒の種類を間違えない**（`mito_train/models/board_ocr.py` の `hand_head` は 14 スロット固定で、スロット index が piece × side をエンコード）。だから「hand の間違い = 枚数の間違い」だけ。 + +残っている失敗パターンは以下の 2 つ： + +### 1. 枚数の隣接ミス（本命） + +`hand_logits: (B, 14, 19)` の 19-way CrossEntropy を使っているため、**「5 vs 6」の間違いも「5 vs 18」の間違いも同じ loss**。序数情報を完全に捨てている。argmax が隣接クラスで揺れやすいのはこの構造由来。sfen 落ちの支配的要因のはず。 + +対策は `TRAINING_PLAN.md#A` の **regression head**（SmoothL1、推論時 round+clamp）。 + +### 2. 0 バイアス（分布の偏り） + +学習データ（`ultemica/piyoshogi` `ocr_paired` train, n=18,000 SFEN）の分布が極端に 0 に寄っている（2026-07-12 実測）： + +- **ラベル 0（空スロット）が 65.68%**。常に 0 を返すだけで slot_acc がこの水準。 +- ラベル 1 が 21.20%、2 が 6.39%、3 以降は 3% 未満に急落 +- 駒種別の非ゼロ率：**飛 17〜18% / 香 22% / 桂 29% / 銀 31% / 金 32% / 角 42% / 歩 64〜68%**。飛は 0/1、たまに 2 でほぼ完結 +- 歩スロットだけが count=1〜18 の全域に散らばる（count=10 も 148/145 件、count=18 も 4 件、以前の欠損は解消済み） +- 高枚数（3〜9）は希少で、under-count しやすい + +対症療法として v2 で `class weight sqrt+clip[0.5, 10.0]` を導入済み。現行データの raw weight は count=6 で 10.27、count=10 で 45.3、count=18 で 1658 なので、clip=10 は **count=7 以降を全部同じ重みにキャップ** している状態。希少枚数を細かく当てにいくには clip を上げるか regression 化が必要。本丸は上記 1 の regression 化。 + +### val 分布の穴（評価に注意） + +val（`ocr_paired` val, 2,000 SFEN, 2026-07-12 実測）は train と分布形が違い、hand の改善評価に穴がある： + +- **val の 0 比率が 75.10%**（train は 65.68%）。「常に 0」ベースラインの val slot_acc が **10pt 底上げ**される。sweep で観測している `val/hand/slot_acc ≒ 0.99` はこの底上げ込み。 +- **count=11 以上が val に 1 件も無い**。歩 count=11〜18 は train に計 500 件超あるが val では見えない → **regression head の高枚数改善効果を val slot_acc では評価できない**。 +- count=10 は val 全体で 1 件のみ（S:P）。学習側は 293 件入ったのに、val での的中は事実上測れない。 +- 飛の非ゼロ率は val で 9.4% / 7.2%（train は 17.8% / 17.1%）。val でさらに希少。 + +**含意**： +- hand の改善効果は `val/hand/slot_acc` の全体値だけでなく、per-count recall（特に count ≥ 3）で追う +- 歩高枚数（11〜18）の効果を測るには合成 mini-eval セットが要る、あるいは `test-realistic` 側で実測する +- backbone sweep の val sfen_full が 33pt 開くのは、slot_acc よりも「たまに出る count ≥ 3 を当てられるか」で決まっている可能性が高い + +### 解消済みの構造欠陥（参考） + +以前は歩スロットで **count=10 が学習データに 1 件も無い** ギャップがあり（piyo-hook 側の SFEN 出力バグ由来）、歩 10 枚を毎回 8/9/11 に誤読していた。現在は歩 count=10 局面が S:P で 148 件、G:P で 145 件、count=18 も 4/4 件収録されており、この構造欠陥は解消されている。分布としては依然として裾で希少（count=10 は全体の 0.12%）なので、対策 2（class weight）の対象には残っている。 + +### 診断コマンド + +「どのスロットで、どの枚数を、何枚に間違えたか」を数字で出す： + +- `scripts/inspect/diagnose_hand.py` — val 全体で per-slot accuracy + true count → pred count の confusion matrix +- `scripts/inspect/analyze_hand_failures.py` — 端末別（iPhone XR / iPhone 15 等）で mismatch 局面をダンプ + +sweep 直後の checkpoint に対して回せば、隣接ミス / 0 バイアスのどちらが支配的か切り分けられる。 + +## 画像サイズを大きくしたときの見込み + +現状 sweep は **image_size=224**。9x9 マス識別に対して 1マスあたり `224/9 ≈ 24.9px`。24x24 の画像で「歩 / と / 銀 / 金」を漢字で判別する状態。マス内の細部（成香=杏、成桂=圭 等の点画）が潰れる領域に片足入っている。 + +`docs/backbones.md#何が精度を押し上げるか` に **「image_size の増（224 → 288 → 384）: 一番効きやすい」** と明記されており、経験則としても解像度は cell_acc に直接効く。 + +### 参考: v1 baseline との落差 + +`TRAINING_PLAN.md` の v1 baseline は mobilenet_v3_small を epoch 40 まで回して **sfen_acc = 0.486**（batch 128, lr 6e-4, image 288 系）。今回の sweep 同 backbone ep40 で **sfen 0.388**。**batch と lr が違うので直接比較はできない**が、image_size の効きが大きな要因の 1 つなのは確か。 + +### 224 → 288 / 384 に上げた場合の予想 + +1マスあたり px 数と、cell_acc / sfen_full の伸び幅の見立て： + +| image_size | 1マスあたり px | cell_acc の伸び (mobilenet_v3_small 想定) | sfen_full 期待伸び幅 | +|---:|---:|---:|---:| +| 224 (現状) | ~24.9 | ― (baseline 0.966) | ― (baseline 0.385) | +| 288 | ~32.0 | +0.5〜1.0pt | **+0.05〜+0.10** | +| 384 | ~42.7 | +0.8〜1.5pt | **+0.10〜+0.20** | + +sfen_full 側の増分が大きいのは、cell_acc の 1pt 改善が **81 マス独立仮定で ~ (0.99/0.98)^81 ≈ 2.3倍** の効きになるため。 + +**image_size を上げる方が、追加エポックを回すよりコストパフォーマンスが良い**。VRAM と学習時間は概ね `(size/224)^2` で増えるが、`convnext_atto` くらいまでのモデルなら 288 / 384 は現実的。 + +### backbone × image_size の組み合わせ予想 (ep50, cosine lr 併用時) + +epoch 50 相当で回したときの val sfen_full の見込みレンジ： + +| Backbone | 224 (実測) | 288 予想 | 384 予想 | +|---|---:|---:|---:| +| mobilenet_v3_small | 0.385 | 0.45〜0.50 | 0.55〜0.60 | +| mobilenet_v3_large | 0.593 | 0.65〜0.70 | 0.72〜0.78 | +| convnext_atto | 0.611 | 0.67〜0.72 | 0.74〜0.80 | +| efficientnet_b0 | 0.661 | 0.72〜0.76 | 0.78〜0.83 | +| convnext_femto | 0.645 | 0.70〜0.75 | 0.77〜0.82 | +| efficientnet_b1 | 0.663 | 0.72〜0.77 | 0.79〜0.84 | +| convnext_pico | 0.634 | 0.71〜0.75 | 0.77〜0.82 | +| convnext_nano | 0.660 | 0.72〜0.76 | 0.78〜0.83 | +| **convnext_tiny** | **0.716** | **0.77〜0.82** | **0.83〜0.88** | + +この見込みは **image_size 効果 (backbones.md の想定)** + **既にある epoch トラジェクトリ** + **hand head の残り改善** を足したもの。convnext_tiny × 384 で **0.85 台**、実用ライン 90% には **もう一押し必要** な位置。 + +## 「エポック増」と「画像サイズ増」の効き順ランキング + +sweep 結果から見えるコスパ順位（1 sample あたりの学習 wall-clock 増加を考慮）： + +1. **image_size 224 → 288**（コスト ~1.65倍、期待 +5〜10pt） +2. **hand head の regression 化 + class weight 継続**（コスト実装のみ、期待 +3〜5pt） +3. **image_size 288 → 384**（コスト ~1.78倍、期待 +5〜10pt） +4. **cosine LR + epoch 100 まで延長**（コスト 2倍、期待 +2〜3pt） +5. **backbone を 1〜2 段上げる**（コストは backbone 依存、期待 +2〜5pt） + +**エポック単独増は 4 番目**。sweep の各 backbone の後半勾配を見る限り、`convnext_atto` `convnext_pico` は既に飽和気味で、100 epoch まで回しても sfen +2pt がせいぜい。 + +一方 **画像サイズ増と hand head 改修が上位で、これを先に済ませないと後段の投資が全部 5割引き** になる。hand が枚数の隣接ミスで落ちている限り、cell_acc を上げても sfen_full には抜けない。 + +## 90% ライン到達までの推奨経路 + +`ocr-metrics.md` の **プライマリ目標 = SFEN Exact-Match 90%** に到達するための現実的な階段： + +``` +[現状 baseline] +convnext_tiny × 224 × 50ep → 0.716 + + ↓ image_size を 288 に (+cosine LR) +convnext_tiny × 288 × 60ep → 0.77 前後 (2026-07 内で到達目安) + + ↓ image_size を 384 に + hand regression head +convnext_tiny × 384 × 80ep → 0.83〜0.86 (2026-08 内で到達目安) + + ↓ 蒸留 (convnext_small teacher) + 追加 aug +convnext_tiny × 384 × distill → 0.88〜0.92 (実用ライン到達) + + ↓ SNS aug ハード側の追加 + 手番/持ち駒後処理 +0.93+ (SNS 実運用ライン) +``` + +**convnext_tiny × 384 に到達しても素の sfen_full は 0.85 前後で、90% には蒸留 or 後処理が要る** というのが今回の sweep から出る現実的な見立て。ブラウザ配信の本命 convnext_nano は 15.1M ながら val 0.660 に留まり、efficientnet_b0 (4.4M) と同点。Pareto フロント上ではむしろ **efficientnet_b0 の方がサイズ効率で優位**、精度上限を狙うなら **convnext_tiny (28.0M)** が唯一の選択肢という構図。 + +## 直近の推奨アクション + +sweep 結果を踏まえた次のイテレーション優先順位： + +1. **hand head の regression 化**（`TRAINING_PLAN.md#A`）。cell_acc は据置きで sfen_full が跳ねる可能性大。実装コスト最小、まず切って効果測定。convnext_nano/pico/atto/femto がまとめて hand 側で律速されている疑いを直接叩ける。 +2. **image_size=288 での再 sweep**。まずは mobilenet_v3_small / convnext_atto / convnext_tiny / **efficientnet_b0** の 4 点で効果測定。b0 は Pareto 効率が良いので nano より優先。 +3. **convnext_nano の扱いを再検討**。15.1M で 0.660 は Pareto 上で b0 (4.4M, 0.661) に負けており、ブラウザ配信本命の座は再評価が必要。tiny の 15MB 版としての位置付けは要更新。 +4. **cosine annealing の導入**（`TRAINING_PLAN.md#D`）。resume 対応と同時に。 +5. **image_size=384 は convnext_atto / tiny のみ、蒸留と一緒に**。VRAM とのトレードで動く backbone だけ選ぶ。 +6. **hand 分布の再集計**（`scripts/inspect/analyze_hand_distribution.py`）。歩 count=10 収録後の最新分布を確認し、class weight の clip 値を再調整するか判断する。 + +「エポック増だけで頑張る」路線は **ROI が悪い**。sweep の後半勾配は既に落ちている。次の投資は **hand head + 解像度 + scheduler の三点セット**、そこに convnext_tiny を据えるのが最短。 + +## 再現用コマンド + +sweep を同条件で再現する場合： + +```bash +# 単一 GPU で全 backbone を順番に +EPOCHS=50 IMAGE_SIZE=224 ./scripts/train_backbones.sh + +# image_size=288 版 +EPOCHS=50 IMAGE_SIZE=288 BACKBONES="mobilenet_v3_small convnext_atto convnext_tiny" \ + ./scripts/train_backbones.sh + +# resume-friendly な全 backbone sweep +RESUME_INCOMPLETE=1 EPOCHS=100 IMAGE_SIZE=288 ./scripts/train_backbones.sh +``` + +W&B project: `mito-train-board-ocr`。今回の sweep run 一覧： + +| Backbone | Run ID | +|---|---| +| mobilenet_v3_small | vhhpiwgm | +| mobilenet_v3_large | k2mol6pl | +| convnext_atto | o68bjpih | +| convnext_femto | sw2z4m7f | +| efficientnet_b0 | 3p1kubgl | +| convnext_pico | vw4mk54g | +| efficientnet_b1 | 5h2kp29v | +| convnext_tiny | r4c7icav | +| convnext_nano | nyrsx25c | diff --git a/mito_train/datasets/capture_dataset.py b/mito_train/datasets/capture_dataset.py index 223d24b..3c0a956 100644 --- a/mito_train/datasets/capture_dataset.py +++ b/mito_train/datasets/capture_dataset.py @@ -50,16 +50,20 @@ def build_transform( return A.Compose([ # Fit longest side to image_size, pad shorter side to square with black bars. # This preserves the entire board + both piece stands without clipping. - A.LongestMaxSize(max_size=image_size), + # INTER_AREA: fastest + best quality for downscaling (sources are + # larger than image_size), vs. the default INTER_LINEAR. + A.LongestMaxSize(max_size=image_size, interpolation=cv2.INTER_AREA), A.PadIfNeeded( min_height=image_size, min_width=image_size, border_mode=cv2.BORDER_CONSTANT, fill=0, ), - # Simulate position/scale jitter within the padded canvas (no content loss). + # Position/scale jitter modeling residual detector error only — the + # upstream board detector normalizes framing, so heavy jitter would + # waste capacity. Kept small to buffer a few px of crop drift. A.Affine( - translate_percent=(-0.05, 0.05), - scale=(0.9, 1.0), - border_mode=cv2.BORDER_CONSTANT, fill=0, p=0.7, + translate_percent=(-0.02, 0.02), + scale=(0.97, 1.03), + border_mode=cv2.BORDER_CONSTANT, fill=0, p=0.5, ), # Color / brightness / saturation (device differences) A.RandomBrightnessContrast( @@ -68,14 +72,18 @@ def build_transform( hue_shift_limit=5, sat_shift_limit=10, val_shift_limit=5, p=0.3), # Noise / degradation A.GaussNoise(std_range=(0.02, 0.08), p=0.3), - A.ImageCompression(quality_range=(50, 90), p=0.7), - A.Downscale(scale_range=(0.5, 0.9), p=0.3), + # Typical single-hop SNS re-encoding (Twitter/Insta/LINE, quality 65-92). + # JPEG encode/decode is CPU-heavy; p tuned to not starve GPU. + A.ImageCompression(quality_range=(65, 92), p=0.4), + # Rare heavy degradation: reshared / multi-hop / aggressive compressors. + A.ImageCompression(quality_range=(40, 60), p=0.05), + A.Downscale(scale_range=(0.5, 0.9), p=0.2), A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), ToTensorV2(), ]) else: return A.Compose([ - A.LongestMaxSize(max_size=image_size), + A.LongestMaxSize(max_size=image_size, interpolation=cv2.INTER_AREA), A.PadIfNeeded( min_height=image_size, min_width=image_size, border_mode=cv2.BORDER_CONSTANT, fill=0, diff --git a/mito_train/datasets/hf_capture_dataset.py b/mito_train/datasets/hf_capture_dataset.py index 5079ad3..5249512 100644 --- a/mito_train/datasets/hf_capture_dataset.py +++ b/mito_train/datasets/hf_capture_dataset.py @@ -1,9 +1,10 @@ """HF Hub-backed variant of CaptureDataset. Loads the dataset from a Hugging Face dataset repo where each row is -{"image": , "sfen": , "hash": } — the vertical (1 row per -image) schema. For the current paired schema (1 row per SFEN, 4 images), use -HFPairedDataset instead. +{"images": [, ...], "devices": [, ...], "sfen": , "hash": }. +Each sfen is paired with multiple device renderings; __getitem__ picks one +device per call (random for train, fixed index for val). See HFPairedDataset +for the alternative flatten-per-device approach. Usage (in training): from mito_train.datasets import HFCaptureDataset, build_transform @@ -17,18 +18,52 @@ Subsequent runs are cache-hits. """ from __future__ import annotations - +import fcntl +import os +import random +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path from typing import Callable import albumentations as A +import cv2 import numpy as np import torch from PIL import Image from torch.utils.data import Dataset +from tqdm.auto import tqdm from .capture_dataset import build_transform from .sfen_utils import parse_sfen +DEFAULT_PRELOAD_CACHE_DIR = Path( + os.environ.get("MITO_PRELOAD_CACHE", Path.home() / ".cache" / "mito-train" / "preload") +) + + +def _decode_resize_pad(entry_bytes: bytes, image_size: int) -> np.ndarray: + """WebP decode + LongestMaxSize + center-pad-to-square (all deterministic). + + Returns (image_size, image_size, 3) uint8 RGB. This encapsulates the fixed + part of the training pipeline so it can be cached once and skipped every + epoch. cv2 ops release the GIL, so a ThreadPoolExecutor scales. + """ + buf = np.frombuffer(entry_bytes, dtype=np.uint8) + img = cv2.imdecode(buf, cv2.IMREAD_COLOR_RGB) + h, w = img.shape[:2] + scale = image_size / max(h, w) + new_h = int(round(h * scale)) + new_w = int(round(w * scale)) + resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA) + pad_h = image_size - new_h + pad_w = image_size - new_w + top = pad_h // 2 + left = pad_w // 2 + return cv2.copyMakeBorder( + resized, top, pad_h - top, left, pad_w - left, + cv2.BORDER_CONSTANT, value=0, + ) + class HFCaptureDataset(Dataset): """CaptureDataset backed by a Hugging Face Hub dataset repo.""" @@ -41,29 +76,60 @@ def __init__( limit: int | None = None, cache_dir: str | None = None, streaming: bool = False, + config_name: str = "ocr_paired", + device_index: int | None = None, + preload: bool = False, + preload_image_size: int = 224, + preload_workers: int = 16, + preload_cache_dir: str | Path | None = None, ) -> None: """ Parameters ---------- - repo_id: Hugging Face dataset repo (e.g. "ultemica/piyoshogi") - split: "train" or "val" - transform: Albumentations Compose or callable. Defaults to val transform. - limit: Truncate to first N rows (smoke tests). - cache_dir: Override HF cache location. None = ~/.cache/huggingface/datasets - streaming: If True, stream from Hub (no local cache). Random access breaks; - __getitem__ becomes O(n). Use only for one-shot iteration. + repo_id: Hugging Face dataset repo (e.g. "ultemica/piyoshogi") + split: "train" or "val" + transform: Albumentations Compose or callable. Defaults to val transform. + limit: Truncate to first N rows (smoke tests). + cache_dir: Override HF cache location. None = ~/.cache/huggingface/datasets + streaming: If True, stream from Hub (no local cache). Random access breaks; + __getitem__ becomes O(n). Use only for one-shot iteration. + config_name: HF dataset config (e.g. "ocr_paired", "detector_paired"). + device_index: Each row holds a list of device renderings for one sfen. + None (default) picks one uniformly at random per __getitem__ + call — good for training (sees all devices across epochs). + int uses that index (clamped) — good for stable val metrics. + preload: If True, decode + resize + pad every image at init time and + hold them in a single uint8 ndarray. Removes WebP decode + and the deterministic resize/pad from the hot loop entirely. + Cost: ~preload_image_size^2 * 3 * N * num_devices bytes of RAM + (~10.5 GB for 18k rows x 4 devices at 224). Fork+COW shares + this across DataLoader workers on Linux. + preload_image_size: Target square size for the preloaded arrays. Should match + the image_size used when building the augmentation transform. + preload_workers: Threads used to decode during preload. cv2.imdecode releases + the GIL so this scales well up to physical core count. + preload_cache_dir: Directory holding the on-disk preload cache. Second and + later runs mmap the cache and skip decode entirely. + None -> $MITO_PRELOAD_CACHE or ~/.cache/mito-train/preload. """ - from datasets import load_dataset + from datasets import Image as HFImage + from datasets import Sequence, load_dataset self.transform = transform if transform is not None else build_transform("val") self.streaming = streaming + self.device_index = device_index ds = load_dataset( repo_id, + config_name, split=split, cache_dir=cache_dir, streaming=streaming, ) + # Disable auto-decode so row access returns raw {bytes, path} per image. + # We pick one device index first, then decode only that single WebP. + # Cuts per-row decode cost by ~len(devices) (typically 4x). + ds = ds.cast_column("images", Sequence(HFImage(decode=False))) if limit is not None and not streaming: ds = ds.select(range(min(limit, len(ds)))) self.ds = ds @@ -74,6 +140,127 @@ def __init__( else: self._length = len(ds) + self._cache: np.ndarray | None = None + if preload and not streaming: + cache_dir = Path(preload_cache_dir) if preload_cache_dir is not None else DEFAULT_PRELOAD_CACHE_DIR + self._preload(preload_image_size, preload_workers, cache_dir, repo_id, config_name, split) + + def _cache_path( + self, cache_dir: Path, repo_id: str, config_name: str, + split: str, image_size: int, + ) -> Path: + """Build a stable cache path keyed by the dataset identity + image size. + + HF's `_fingerprint` also encodes any `.select(...)` truncation, so a + `--limit N` run has a different cache from the full run — no risk of + pulling a truncated cache in full mode. + """ + fingerprint = getattr(self.ds, "_fingerprint", "nofp") + safe_repo = repo_id.replace("/", "_") + name = f"{safe_repo}__{config_name}__{split}__{fingerprint}__sz{image_size}.npy" + return cache_dir / name + + def _preload( + self, image_size: int, workers: int, cache_dir: Path, + repo_id: str, config_name: str, split: str, + ) -> None: + """Decode + resize + pad every device rendering into an on-disk ndarray. + + Shape: (N, num_devices, image_size, image_size, 3) uint8 RGB. + + First call writes to `cache_dir/.npy`; subsequent calls memory-map + the file and skip decode entirely (~30-60s startup collapses to <1s). + + DDP: only rank 0 builds. Other ranks wait on a barrier and mmap the same + file, sharing the kernel page cache instead of duplicating the array. + + PARALLEL_GPU sweeps: many independent Python processes may race on the + same cache path. An fcntl advisory lock on a sibling .lock file + serializes them cross-process — the first process builds, the rest wait + on the lock and then find the cache already there. + """ + from mito_train.training.dist_utils import barrier, is_main + + path = self._cache_path(cache_dir, repo_id, config_name, split, image_size) + + if is_main(): + self._maybe_build_with_lock(path, image_size, workers) + + # All ranks meet here: rank 0 has finished writing, others were idle. + barrier() + + arr = np.load(path, mmap_mode="r") + gb = arr.nbytes / (1024 ** 3) + if is_main(): + print(f"[HFCaptureDataset] preload cache ready: {path} ({gb:.2f} GB, mmap)") + self._cache = arr + + def _maybe_build_with_lock( + self, path: Path, image_size: int, workers: int, + ) -> None: + """Cross-process safe cache-or-build. + + Fast path: file exists -> return immediately without touching the lock. + Slow path: take an exclusive fcntl lock on a sibling .lock file, then + re-check under lock (another process may have built while we waited) + before doing the actual decode + save. + """ + if path.exists(): + return + + path.parent.mkdir(parents=True, exist_ok=True) + lock_path = path.with_name(path.name + ".lock") + + with open(lock_path, "w") as lockf: + print(f"[HFCaptureDataset] acquiring build lock: {lock_path}") + fcntl.flock(lockf.fileno(), fcntl.LOCK_EX) + # Re-check: while we were blocked, another process may have finished. + if path.exists(): + print("[HFCaptureDataset] cache built by another process while waiting; skipping") + return + self._build_and_save_cache(path, image_size, workers) + + def _build_and_save_cache( + self, path: Path, image_size: int, workers: int, + ) -> None: + """Decode all rows and persist as an atomic .npy file. Rank 0 only.""" + n = self._length + first_row = self.ds[0] + num_devices = len(first_row["images"]) + cache = np.zeros( + (n, num_devices, image_size, image_size, 3), dtype=np.uint8, + ) + + # Grab the raw bytes column up-front so worker threads don't hit the + # HF dataset object concurrently (its indexing isn't thread-safe). + all_images = self.ds["images"] + + def process(idx: int) -> None: + entries = all_images[idx] + for dev_idx in range(min(num_devices, len(entries))): + cache[idx, dev_idx] = _decode_resize_pad( + entries[dev_idx]["bytes"], image_size, + ) + for dev_idx in range(len(entries), num_devices): + cache[idx, dev_idx] = cache[idx, len(entries) - 1] + + with ThreadPoolExecutor(max_workers=workers) as pool: + list(tqdm( + pool.map(process, range(n)), + total=n, desc=f"preload {image_size}px", + )) + gb = cache.nbytes / (1024 ** 3) + print(f"[HFCaptureDataset] preloaded {n} rows x {num_devices} devices = {gb:.2f} GB") + + # Atomic write: tmp then rename. Handles Ctrl-C mid-write cleanly. + # File-object form so numpy doesn't re-append ".npy" to our tmp name. + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "wb") as f: + np.save(f, cache) + os.replace(tmp, path) + print(f"[HFCaptureDataset] preload cache saved: {path}") + @property def entries(self): """Sfen-only entries list, compatible with compute_hand_class_weights.""" @@ -93,18 +280,36 @@ def __len__(self) -> int: def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if self.streaming: raise RuntimeError("Streaming datasets do not support random access; iterate instead.") - row = self.ds[idx] - img = row["image"] # PIL Image (decoded by HF Image feature) - - if img.mode == "RGBA": - bg = Image.new("RGB", img.size, (255, 255, 255)) - bg.paste(img, mask=img.split()[3]) - img = bg - elif img.mode != "RGB": - img = img.convert("RGB") - img_np = np.array(img) - - parsed = parse_sfen(row["sfen"]) + + if self._cache is not None: + # Preloaded path: decode + resize + pad were already done at init. + # LongestMaxSize and PadIfNeeded in the transform become no-ops + # because the input is already the target square size. + num_devices = self._cache.shape[1] + if self.device_index is None: + dev_idx = random.randrange(num_devices) + else: + dev_idx = min(self.device_index, num_devices - 1) + # .copy() defends against Albumentations in-place ops mutating the cache. + img_np = self._cache[idx, dev_idx].copy() + sfen = self.ds[idx]["sfen"] + else: + row = self.ds[idx] + images = row["images"] # list[{"bytes": ..., "path": ...}] — undecoded + if self.device_index is None: + entry = random.choice(images) + else: + entry = images[min(self.device_index, len(images) - 1)] + + # cv2.imdecode releases the GIL (C++ libwebp), so multiple workers + # actually decode in parallel. PIL holds the GIL during decode. + # IMREAD_COLOR_RGB (OpenCV >= 4.10) decodes straight to RGB, saving + # the extra full-image cvtColor copy per sample. + buf = np.frombuffer(entry["bytes"], dtype=np.uint8) + img_np = cv2.imdecode(buf, cv2.IMREAD_COLOR_RGB) + sfen = row["sfen"] + + parsed = parse_sfen(sfen) board = torch.tensor(parsed.board, dtype=torch.long) # (9,9) hand = torch.tensor(parsed.hand, dtype=torch.long) # (14,) @@ -112,6 +317,6 @@ def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tenso out = self.transform(image=img_np) image_tensor = out["image"] else: - image_tensor = self.transform(img) + image_tensor = self.transform(Image.fromarray(img_np)) return image_tensor, board, hand diff --git a/mito_train/models/board_ocr.py b/mito_train/models/board_ocr.py index eb28d2a..1d0b471 100644 --- a/mito_train/models/board_ocr.py +++ b/mito_train/models/board_ocr.py @@ -19,32 +19,87 @@ import torch import torch.nn as nn -BackboneName = Literal["mobilenet_v3_small", "convnext_tiny"] +BackboneName = Literal[ + "mobilenet_v3_small", + "mobilenet_v3_large", + "efficientnet_b0", + "efficientnet_b1", + "convnext_atto", + "convnext_femto", + "convnext_pico", + "convnext_nano", + "convnext_tiny", +] HandMode = Literal["classification", "regression"] +# timm's convnext family (atto/femto/pico/nano) — same architecture as +# torchvision's convnext_tiny, scaled down. Loaded uniformly via _TimmFeatureWrapper. +_TIMM_CONVNEXT_VARIANTS: tuple[str, ...] = ( + "convnext_atto", + "convnext_femto", + "convnext_pico", + "convnext_nano", +) + # Theoretical max count per hand slot (sente P,L,N,S,G,B,R then gote same order). # Pawn = 18, minor pieces = 4, bishop/rook = 2. PIECE_MAX_PER_SLOT: tuple[int, ...] = (18, 4, 4, 4, 4, 2, 2, 18, 4, 4, 4, 4, 2, 2) def _build_backbone(name: BackboneName, pretrained: bool) -> tuple[nn.Module, int]: - """Return (features_module, out_channels).""" + """Return (features_module, out_channels). + + torchvision backbones expose feature maps directly via `.features`. + timm backbones use `features_only=True` and return a list of feature + maps at each stride; we take the deepest (last) one. + """ if name == "mobilenet_v3_small": - from torchvision.models import ( - MobileNet_V3_Small_Weights, - mobilenet_v3_small, - ) + from torchvision.models import MobileNet_V3_Small_Weights, mobilenet_v3_small weights = MobileNet_V3_Small_Weights.DEFAULT if pretrained else None m = mobilenet_v3_small(weights=weights) return m.features, 576 + if name == "mobilenet_v3_large": + from torchvision.models import MobileNet_V3_Large_Weights, mobilenet_v3_large + weights = MobileNet_V3_Large_Weights.DEFAULT if pretrained else None + m = mobilenet_v3_large(weights=weights) + return m.features, 960 + if name == "efficientnet_b0": + from torchvision.models import EfficientNet_B0_Weights, efficientnet_b0 + weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None + m = efficientnet_b0(weights=weights) + return m.features, 1280 + if name == "efficientnet_b1": + from torchvision.models import EfficientNet_B1_Weights, efficientnet_b1 + weights = EfficientNet_B1_Weights.DEFAULT if pretrained else None + m = efficientnet_b1(weights=weights) + return m.features, 1280 if name == "convnext_tiny": from torchvision.models import ConvNeXt_Tiny_Weights, convnext_tiny weights = ConvNeXt_Tiny_Weights.DEFAULT if pretrained else None m = convnext_tiny(weights=weights) return m.features, 768 + if name in _TIMM_CONVNEXT_VARIANTS: + import timm + m = timm.create_model(name, pretrained=pretrained, features_only=True) + feat_dim = m.feature_info.channels()[-1] + return _TimmFeatureWrapper(m), feat_dim raise ValueError(f"unknown backbone: {name}") +class _TimmFeatureWrapper(nn.Module): + """Adapts timm's features_only output (list of feature maps) to expose only + the deepest map, matching the torchvision `.features` contract used elsewhere. + """ + + def __init__(self, backbone: nn.Module) -> None: + super().__init__() + self.backbone = backbone + + def forward(self, x: torch.Tensor) -> torch.Tensor: + feats = self.backbone(x) + return feats[-1] + + class BoardOCR(nn.Module): """Two-head model: capture image -> (board grid logits, hand slot counts). diff --git a/mito_train/training/dist_utils.py b/mito_train/training/dist_utils.py new file mode 100644 index 0000000..d8ca311 --- /dev/null +++ b/mito_train/training/dist_utils.py @@ -0,0 +1,95 @@ +"""Small helpers around torch.distributed so single-GPU and DDP share the same +training script. + +Usage pattern (see train_board_ocr.py): + setup_distributed() # no-op unless launched via torchrun + if is_main(): # log/save/wandb only on rank 0 + ... + dataset ... DistributedSampler(...) + model ... DDP(model, device_ids=[local_rank]) + metrics ... all_reduce_mean(tensor) # aggregate across ranks + teardown_distributed() + +Detection is env-driven: torchrun exports LOCAL_RANK / RANK / WORLD_SIZE. +When those aren't set the helpers degrade to single-process behavior, so the +same code path runs on the 1-GPU dev box and the 8-GPU A100 node.""" +from __future__ import annotations + +import os + +import torch +import torch.distributed as dist + + +def is_distributed() -> bool: + return int(os.environ.get("WORLD_SIZE", "1")) > 1 + + +def get_world_size() -> int: + return int(os.environ.get("WORLD_SIZE", "1")) + + +def get_rank() -> int: + return int(os.environ.get("RANK", "0")) + + +def get_local_rank() -> int: + return int(os.environ.get("LOCAL_RANK", "0")) + + +def is_main() -> bool: + return get_rank() == 0 + + +def setup_distributed() -> None: + """Initialize the NCCL process group and bind this process to its GPU. + + Safe to call unconditionally; it's a no-op unless torchrun set WORLD_SIZE. + """ + if not is_distributed(): + return + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + if torch.cuda.is_available(): + torch.cuda.set_device(get_local_rank()) + + +def teardown_distributed() -> None: + if is_distributed() and dist.is_initialized(): + dist.destroy_process_group() + + +def barrier() -> None: + """Rendezvous point across ranks. No-op outside DDP.""" + if is_distributed() and dist.is_initialized(): + dist.barrier() + + +def all_reduce_mean(tensor: torch.Tensor) -> torch.Tensor: + """Average a scalar tensor across ranks (in place). No-op outside DDP. + + Use this to aggregate per-rank running metrics into a global mean at + epoch end; each rank saw len(dataset)/world_size samples so a plain + mean-of-per-rank-means is unweighted-correct only when the sampler + hands out equal-size shards, which DistributedSampler does by design + (it drops or pads to make shards equal). + """ + if is_distributed() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + tensor = tensor / get_world_size() + return tensor + + +def resolve_device() -> str: + """Return the device string this rank should train on. + + Under DDP each process pins to its LOCAL_RANK GPU; outside DDP fall back + to whatever accelerator the harness's get_device() picks. + """ + if is_distributed() and torch.cuda.is_available(): + return f"cuda:{get_local_rank()}" + if torch.cuda.is_available(): + return "cuda" + if torch.backends.mps.is_available(): + return "mps" + return "cpu" diff --git a/mito_train/training/train_board_ocr.py b/mito_train/training/train_board_ocr.py index 4a07e00..6287f47 100644 --- a/mito_train/training/train_board_ocr.py +++ b/mito_train/training/train_board_ocr.py @@ -1,10 +1,18 @@ """Training entry point for BoardOCR (image -> 9x9 board grid + 14-slot hand counts). -Currently only smoke mode is available: +Single GPU (smoke): python -m mito_train.training.train_board_ocr --mode smoke --limit 100 --epochs 5 -Real training (intended to run on a GPU machine) uses --mode full to iterate over all 27k samples: +Single GPU (full): python -m mito_train.training.train_board_ocr --mode full --epochs 20 + +Multi-GPU (single node, N GPUs) via torchrun: + torchrun --standalone --nproc_per_node=N \\ + -m mito_train.training.train_board_ocr --mode full --epochs 20 + +The same script handles both launch modes: `setup_distributed()` no-ops outside +torchrun, so rank/world detection, DDP wrapping, sampler shard rotation, and +metric all_reduce all activate only when actually distributed. """ from __future__ import annotations @@ -12,13 +20,39 @@ from pathlib import Path import torch +import torch.nn as nn import torch.nn.functional as F -from torch.utils.data import ConcatDataset, DataLoader +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from tqdm.auto import tqdm from mito_train.datasets import CaptureDataset, build_transform from mito_train.datasets.sfen_utils import parse_sfen from mito_train.models import BoardOCR -from mito_train.training.train_piece import _init_wandb, get_device +from mito_train.training.dist_utils import ( + all_reduce_mean, + get_local_rank, + get_rank, + get_world_size, + is_distributed, + is_main, + resolve_device, + setup_distributed, + teardown_distributed, +) +from mito_train.training.train_piece import _init_wandb + + +def _unwrap(model: nn.Module) -> nn.Module: + """Strip DDP and torch.compile wrappers to reach the underlying module. + + Needed when saving/loading state_dict: DDP prefixes with `module.` and + compile prefixes with `_orig_mod.` — unwrapping both makes checkpoints + portable across launch modes. + """ + m = model.module if isinstance(model, DDP) else model + return getattr(m, "_orig_mod", m) def compute_hand_class_weights( @@ -47,51 +81,57 @@ def compute_hand_class_weights( def compute_loss( board_logits: torch.Tensor, # (B, 29, 9, 9) - hand_out: torch.Tensor, # (B, 14, 19) or (B, 14) + hand_logits: torch.Tensor, # (B, 14, 19) board_target: torch.Tensor, # (B, 9, 9) hand_target: torch.Tensor, # (B, 14) - hand_mode: str = "classification", hand_weight: float = 1.0, hand_class_weight: torch.Tensor | None = None, - smooth_l1_beta: float = 1.0, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Weighted sum of board CE + hand loss (CE for classification / SmoothL1 for regression).""" + """Weighted sum of board CE + hand CE.""" board_loss = F.cross_entropy(board_logits, board_target) - if hand_mode == "classification": - B, S, C = hand_out.shape - hand_loss = F.cross_entropy( - hand_out.reshape(B * S, C), - hand_target.reshape(B * S), - weight=hand_class_weight, - ) - elif hand_mode == "regression": - hand_loss = F.smooth_l1_loss(hand_out, hand_target.float(), beta=smooth_l1_beta) - else: - raise ValueError(f"unknown hand_mode: {hand_mode!r}") + B, S, C = hand_logits.shape + hand_loss = F.cross_entropy( + hand_logits.reshape(B * S, C), + hand_target.reshape(B * S), + weight=hand_class_weight, + ) total = board_loss + hand_weight * hand_loss return total, board_loss, hand_loss +METRIC_KEYS = ( + "board/cell_acc", "board/full_acc", + "hand/slot_acc", "hand/full_acc", "sfen/full_acc", +) + + def compute_metrics( board_logits: torch.Tensor, - hand_out: torch.Tensor, + hand_logits: torch.Tensor, board_target: torch.Tensor, hand_target: torch.Tensor, - hand_pred: torch.Tensor, # (B, 14) discrete counts from model.predict_hand -) -> dict[str, float]: +) -> dict[str, torch.Tensor]: + """Per-batch accuracy metrics as 0-dim GPU tensors. + + Returns tensors (not floats) so the hot loop can accumulate on-device + without forcing a CPU<->GPU sync every step; call .item() only when + actually logging/printing. + """ with torch.no_grad(): board_pred = board_logits.argmax(dim=1) # (B, 9, 9) - board_cell_acc = (board_pred == board_target).float().mean().item() + board_ok = board_pred == board_target + board_cell_acc = board_ok.float().mean() # Fraction of images where all 81 cells are correct - board_full = (board_pred == board_target).all(dim=(1, 2)).float().mean().item() + board_all = board_ok.all(dim=(1, 2)) + board_full = board_all.float().mean() - hand_slot_acc = (hand_pred == hand_target).float().mean().item() - hand_full = (hand_pred == hand_target).all(dim=1).float().mean().item() + hand_pred = hand_logits.argmax(dim=-1) # (B, 14) + hand_ok = hand_pred == hand_target + hand_slot_acc = hand_ok.float().mean() + hand_all = hand_ok.all(dim=1) + hand_full = hand_all.float().mean() - sfen_full = ( - ((board_pred == board_target).all(dim=(1, 2))) - & ((hand_pred == hand_target).all(dim=1)) - ).float().mean().item() + sfen_full = (board_all & hand_all).float().mean() return { "board/cell_acc": board_cell_acc, "board/full_acc": board_full, @@ -102,8 +142,37 @@ def compute_metrics( def run(args: argparse.Namespace) -> None: - device = get_device() - print(f"[board_ocr] device={device} mode={args.mode} backbone={args.backbone}") + # torchrun sets WORLD_SIZE/RANK/LOCAL_RANK; single-process runs skip this cleanly. + setup_distributed() + device = resolve_device() + use_cuda = device.startswith("cuda") + world_size = get_world_size() + ddp = is_distributed() + + def log_main(msg: str) -> None: + if is_main(): + print(msg) + + if use_cuda: + # Input size is fixed (image_size x image_size) -> let cuDNN autotune kernels. + torch.backends.cudnn.benchmark = True + # TF32 matmuls on Ampere+: big speedup, negligible precision impact here. + torch.set_float32_matmul_precision("high") + + # Mixed precision: bf16 on Ampere+ (no scaler needed), fp16 + GradScaler otherwise. + amp_dtype = ( + torch.bfloat16 + if use_cuda and torch.cuda.is_bf16_supported() + else torch.float16 + ) + scaler = torch.amp.GradScaler( + "cuda", enabled=use_cuda and amp_dtype == torch.float16, + ) + log_main( + f"[board_ocr] device={device} mode={args.mode} backbone={args.backbone} " + f"amp={amp_dtype if use_cuda else 'off'} " + f"world_size={world_size} rank={get_rank()}" + ) train_tf = build_transform("train", image_size=args.image_size) val_tf = build_transform("val", image_size=args.image_size) @@ -112,133 +181,126 @@ def run(args: argparse.Namespace) -> None: limit_val = min(args.limit, 20) if args.mode == "smoke" else None if args.hf_repo_id: - # HF-hosted parquet dataset. First run downloads shards into HF cache - # (~32GB for ocr_paired); subsequent runs are cache-hits. - if args.hf_config.endswith("_paired"): - from mito_train.datasets import HFPairedDataset - print(f"[board_ocr] loading from HF repo: {args.hf_repo_id} " - f"(config={args.hf_config}, devices={args.hf_devices or 'all'})") - train_ds = HFPairedDataset( - repo_id=args.hf_repo_id, split="train", config=args.hf_config, - transform=train_tf, devices=args.hf_devices or None, - limit=limit_train, - ) - val_ds = HFPairedDataset( - repo_id=args.hf_repo_id, split="val", config=args.hf_config, - transform=val_tf, devices=args.hf_devices or None, - limit=limit_val, - ) - else: - from mito_train.datasets import HFCaptureDataset - print(f"[board_ocr] loading from HF repo: {args.hf_repo_id} (vertical config)") - train_ds = HFCaptureDataset( - repo_id=args.hf_repo_id, split="train", - transform=train_tf, limit=limit_train, - ) - val_ds = HFCaptureDataset( - repo_id=args.hf_repo_id, split="val", - transform=val_tf, limit=limit_val, - ) + # HF-hosted parquet dataset (embedded PNG bytes). First run downloads + # the shards into HF cache (~21GB); subsequent runs are cache-hits. + from mito_train.datasets import HFCaptureDataset + log_main(f"[board_ocr] loading from HF repo: {args.hf_repo_id}") + preload_kwargs = dict( + preload=args.preload, + preload_image_size=args.image_size, + preload_workers=max(args.num_workers, 8), + ) + train_ds = HFCaptureDataset( + repo_id=args.hf_repo_id, split="train", + transform=train_tf, limit=limit_train, + **preload_kwargs, + ) + val_ds = HFCaptureDataset( + repo_id=args.hf_repo_id, split="val", + transform=val_tf, limit=limit_val, + device_index=0, # deterministic device pick for stable val curves + **preload_kwargs, + ) else: - # Primary source (backwards compatible with single-device runs). - train_datasets = [CaptureDataset( + train_ds = CaptureDataset( manifest_path=args.train_manifest, image_root=args.image_root, transform=train_tf, limit=limit_train, - )] - val_datasets = [CaptureDataset( + ) + val_ds = CaptureDataset( manifest_path=args.val_manifest, image_root=args.image_root, transform=val_tf, limit=limit_val, - )] - # Extra sources: "manifest.jsonl:image_dir" pairs. Enables per-device - # captures (e.g. iPhone11,8, iPhone15,4 for high-hand augmentation). - for src in args.extra_train_source: - manifest, root = src.split(":", 1) - print(f"[board_ocr] + extra train: {manifest} <- {root}") - train_datasets.append(CaptureDataset( - manifest_path=Path(manifest), image_root=Path(root), - transform=train_tf, limit=limit_train, - )) - for src in args.extra_val_source: - manifest, root = src.split(":", 1) - print(f"[board_ocr] + extra val: {manifest} <- {root}") - val_datasets.append(CaptureDataset( - manifest_path=Path(manifest), image_root=Path(root), - transform=val_tf, limit=limit_val, - )) - train_ds = ConcatDataset(train_datasets) if len(train_datasets) > 1 else train_datasets[0] - val_ds = ConcatDataset(val_datasets) if len(val_datasets) > 1 else val_datasets[0] - print(f"[board_ocr] train={len(train_ds)} val={len(val_ds)}") + ) + log_main(f"[board_ocr] train={len(train_ds)} val={len(val_ds)}") + # DistributedSampler shards the dataset across ranks and re-shuffles per + # epoch (via set_epoch). Its shuffle replaces DataLoader's, so we pass + # shuffle=False when a sampler is used. + train_sampler = ( + DistributedSampler(train_ds, shuffle=True, drop_last=False) if ddp else None + ) + val_sampler = ( + DistributedSampler(val_ds, shuffle=False, drop_last=False) if ddp else None + ) + loader_kwargs = dict( + num_workers=args.num_workers, + pin_memory=use_cuda, + persistent_workers=(args.num_workers > 0), + prefetch_factor=(args.prefetch_factor if args.num_workers > 0 else None), + ) train_loader = DataLoader( - train_ds, batch_size=args.batch_size, shuffle=True, - num_workers=args.num_workers, pin_memory=(device == "cuda"), + train_ds, batch_size=args.batch_size, sampler=train_sampler, + shuffle=(train_sampler is None), **loader_kwargs, ) val_loader = DataLoader( - val_ds, batch_size=args.batch_size, shuffle=False, - num_workers=args.num_workers, pin_memory=(device == "cuda"), + val_ds, batch_size=args.batch_size, sampler=val_sampler, + shuffle=False, **loader_kwargs, ) - model = BoardOCR( - backbone=args.backbone, hand_mode=args.hand_mode, pretrained=args.pretrained, - ).to(device) - n_params = sum(p.numel() for p in model.parameters()) - print(f"[board_ocr] {args.backbone} hand_mode={args.hand_mode} params={n_params:,}") - - if args.hand_mode == "classification" and args.hand_class_weight: - # ConcatDataset doesn't expose .entries directly — merge from children. - if isinstance(train_ds, ConcatDataset): - all_entries = [e for ds in train_ds.datasets for e in ds.entries] - else: - all_entries = train_ds.entries + model = BoardOCR(backbone=args.backbone, pretrained=args.pretrained).to(device) + if use_cuda: + # NHWC layout: faster conv kernels under cuDNN/TensorCores. + model = model.to(memory_format=torch.channels_last) + if ddp: + # Sync BatchNorm stats across ranks so BN-heavy backbones (mobilenet, + # efficientnet) don't diverge on per-rank stats. Harmless for LayerNorm + # backbones (convnext) since no BN layers get replaced. + model = nn.SyncBatchNorm.convert_sync_batchnorm(model) + model = DDP(model, device_ids=[get_local_rank()]) + if args.compile: + model = torch.compile(model) + log_main("[board_ocr] torch.compile enabled (first steps will be slow while compiling)") + n_params = sum(p.numel() for p in _unwrap(model).parameters()) + log_main(f"[board_ocr] {args.backbone} params={n_params:,}") + + if args.hand_class_weight: hand_class_weight = compute_hand_class_weights( - all_entries, - num_classes=model.hand_max, + train_ds.entries, + num_classes=_unwrap(model).hand_max, clip_min=args.class_weight_clip_min, clip_max=args.class_weight_clip_max, ).to(device) - print( + log_main( "[board_ocr] hand class weights (sqrt+clip): " + ", ".join(f"{i}:{w:.2f}" for i, w in enumerate(hand_class_weight.tolist())) ) else: hand_class_weight = None - if args.hand_mode == "regression" and args.hand_class_weight: - print("[board_ocr] --hand-class-weight ignored under hand-mode=regression.") - optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) + # fused=True runs the whole AdamW update in one CUDA kernel per dtype group. + optimizer = torch.optim.AdamW( + model.parameters(), lr=args.lr, weight_decay=1e-4, fused=use_cuda, + ) - args.ckpt_dir.mkdir(parents=True, exist_ok=True) + if is_main(): + args.ckpt_dir.mkdir(parents=True, exist_ok=True) start_epoch = 1 resumed_from: str | None = None resumed_wandb_id: str | None = None if args.resume is not None: ckpt_path = args.ckpt_dir / "latest.pt" if str(args.resume) == "latest" else args.resume - print(f"[board_ocr] resume from {ckpt_path}") + log_main(f"[board_ocr] resume from {ckpt_path}") ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) - ckpt_hand_mode = ckpt.get("hand_mode", "classification") - if ckpt_hand_mode != args.hand_mode: - raise SystemExit( - f"[board_ocr] hand_mode mismatch: ckpt is {ckpt_hand_mode!r} but " - f"--hand-mode={args.hand_mode!r}. Head shapes differ; resume aborted." - ) - model.load_state_dict(ckpt["model"]) + _unwrap(model).load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) + if ckpt.get("scaler") is not None and scaler.is_enabled(): + scaler.load_state_dict(ckpt["scaler"]) start_epoch = ckpt["epoch"] + 1 resumed_from = str(ckpt_path) resumed_wandb_id = ckpt.get("wandb_run_id") - print(f"[board_ocr] resumed at epoch={start_epoch} (ckpt was epoch {ckpt['epoch']})") + log_main(f"[board_ocr] resumed at epoch={start_epoch} (ckpt was epoch {ckpt['epoch']})") if resumed_wandb_id: - print(f"[board_ocr] will resume wandb run id={resumed_wandb_id}") + log_main(f"[board_ocr] will resume wandb run id={resumed_wandb_id}") if args.wandb_run_id: resumed_wandb_id = args.wandb_run_id - print(f"[board_ocr] wandb run id overridden by CLI: {resumed_wandb_id}") + log_main(f"[board_ocr] wandb run id overridden by CLI: {resumed_wandb_id}") - run_name = f"board-ocr-{args.mode}-{args.backbone}" + # Only rank 0 talks to W&B; other ranks keep wandb_run=None and log nothing. + run_name = args.backbone wandb_run = _init_wandb( project="mito-train-board-ocr", run_name=run_name, @@ -250,12 +312,12 @@ def run(args: argparse.Namespace) -> None: "pretrained": args.pretrained, "image_size": args.image_size, "batch_size": args.batch_size, + "effective_batch_size": args.batch_size * world_size, + "world_size": world_size, "lr": args.lr, "epochs": args.epochs, "start_epoch": start_epoch, "resumed_from": resumed_from, - "hand_mode": args.hand_mode, - "smooth_l1_beta": args.smooth_l1_beta, "hand_weight": args.hand_weight, "hand_class_weight": args.hand_class_weight, "class_weight_clip_min": args.class_weight_clip_min, @@ -265,151 +327,216 @@ def run(args: argparse.Namespace) -> None: "n_params": n_params, "device": device, }, - ) + ) if is_main() else None + global_step = 0 for epoch in range(start_epoch, args.epochs + 1): + # DistributedSampler.set_epoch(epoch) rotates the shuffle seed per epoch + # so shards see different orderings; skipping it makes every epoch + # identical inside a single training run. + if train_sampler is not None: + train_sampler.set_epoch(epoch) model.train() - running_loss = 0.0 - running_board_loss = 0.0 - running_hand_loss = 0.0 - running_metrics = { - "board/cell_acc": 0.0, "board/full_acc": 0.0, - "hand/slot_acc": 0.0, "hand/full_acc": 0.0, "sfen/full_acc": 0.0, - } + # Accumulate losses/metrics as 0-dim tensors on-device; .item() forces a + # GPU sync, so we only touch host values every --log-every steps and at + # epoch end. This keeps the CUDA queue full between steps. + running_loss = torch.zeros((), device=device) + running_board_loss = torch.zeros((), device=device) + running_hand_loss = torch.zeros((), device=device) + running_metrics = {k: torch.zeros((), device=device) for k in METRIC_KEYS} n_batches = 0 - for img, board, hand in train_loader: - img = img.to(device) - board = board.to(device) - hand = hand.to(device) - board_logits, hand_out = model(img) - total, bl, hl = compute_loss( - board_logits, hand_out, board, hand, - hand_mode=args.hand_mode, - hand_weight=args.hand_weight, - hand_class_weight=hand_class_weight, - smooth_l1_beta=args.smooth_l1_beta, - ) - optimizer.zero_grad() - total.backward() - optimizer.step() - running_loss += total.item() - running_board_loss += bl.item() - running_hand_loss += hl.item() - hand_pred = model.predict_hand(hand_out) - m = compute_metrics(board_logits, hand_out, board, hand, hand_pred) + pbar = tqdm( + train_loader, desc=f"epoch {epoch:3d}/{args.epochs}", + dynamic_ncols=True, leave=False, disable=not is_main(), + ) + for img, board, hand in pbar: + img = img.to(device, non_blocking=True) + if use_cuda: + img = img.contiguous(memory_format=torch.channels_last) + board = board.to(device, non_blocking=True) + hand = hand.to(device, non_blocking=True) + with torch.autocast("cuda", dtype=amp_dtype, enabled=use_cuda): + board_logits, hand_logits = model(img) + total, bl, hl = compute_loss( + board_logits, hand_logits, board, hand, + hand_weight=args.hand_weight, + hand_class_weight=hand_class_weight, + ) + optimizer.zero_grad(set_to_none=True) + scaler.scale(total).backward() + scaler.step(optimizer) + scaler.update() + running_loss += total.detach() + running_board_loss += bl.detach() + running_hand_loss += hl.detach() + m = compute_metrics(board_logits, hand_logits, board, hand) for k, v in m.items(): running_metrics[k] += v n_batches += 1 + global_step += 1 + if global_step % args.log_every == 0: + # Single sync point: pull host values for tqdm + W&B together. + pbar.set_postfix( + loss=f"{running_loss.item() / n_batches:.3f}", + cell=f"{running_metrics['board/cell_acc'].item() / n_batches:.3f}", + ) + if wandb_run is not None: + wandb_run.log({ + "step/loss": total.item(), + "step/board_loss": bl.item(), + "step/hand_loss": hl.item(), + "step/board/cell_acc": m["board/cell_acc"].item(), + "step/hand/slot_acc": m["hand/slot_acc"].item(), + "step/global_step": global_step, + "step/epoch_frac": epoch - 1 + n_batches / len(train_loader), + }) + + # Reduce per-rank means to a single global mean per metric before + # printing / logging. DistributedSampler shards are equal-size so an + # unweighted mean is unbiased. + avg = { + k: all_reduce_mean(t / n_batches).item() + for k, t in running_metrics.items() + } + avg_loss = all_reduce_mean(running_loss / n_batches).item() + avg_bl = all_reduce_mean(running_board_loss / n_batches).item() + avg_hl = all_reduce_mean(running_hand_loss / n_batches).item() - avg = {k: v / n_batches for k, v in running_metrics.items()} - avg_loss = running_loss / n_batches - avg_bl = running_board_loss / n_batches - avg_hl = running_hand_loss / n_batches - - print( + log_main( f"[board_ocr] epoch={epoch:3d} " f"loss={avg_loss:.4f} (board={avg_bl:.4f} hand={avg_hl:.4f}) " f"cell_acc={avg['board/cell_acc']:.3f} " f"sfen_acc={avg['sfen/full_acc']:.3f}" ) - log = { + wandb_log = { "train/loss": avg_loss, "train/board_loss": avg_bl, "train/hand_loss": avg_hl, - "epoch": epoch, + "train/epoch": epoch, } - log.update({f"train/{k}": v for k, v in avg.items()}) + wandb_log.update({f"train/{k}": v for k, v in avg.items()}) # Simple val (in smoke mode this is only a same-distribution sanity check against train) if epoch % args.val_every == 0 or epoch == args.epochs: model.eval() - val_metrics = {k: 0.0 for k in running_metrics} + val_metrics = {k: torch.zeros((), device=device) for k in METRIC_KEYS} n_val_batches = 0 with torch.no_grad(): - for img, board, hand in val_loader: - img = img.to(device) - board = board.to(device) - hand = hand.to(device) - bl_out, hl_out = model(img) - hand_pred = model.predict_hand(hl_out) - m = compute_metrics(bl_out, hl_out, board, hand, hand_pred) + val_pbar = tqdm( + val_loader, desc=f" val {epoch:3d}", + dynamic_ncols=True, leave=False, disable=not is_main(), + ) + for img, board, hand in val_pbar: + img = img.to(device, non_blocking=True) + if use_cuda: + img = img.contiguous(memory_format=torch.channels_last) + board = board.to(device, non_blocking=True) + hand = hand.to(device, non_blocking=True) + with torch.autocast("cuda", dtype=amp_dtype, enabled=use_cuda): + bl_out, hl_out = model(img) + m = compute_metrics(bl_out, hl_out, board, hand) for k, v in m.items(): val_metrics[k] += v n_val_batches += 1 if n_val_batches > 0: - v_avg = {k: v / n_val_batches for k, v in val_metrics.items()} - print( + v_avg = { + k: all_reduce_mean(t / n_val_batches).item() + for k, t in val_metrics.items() + } + log_main( f"[board_ocr] val cell_acc={v_avg['board/cell_acc']:.3f} " f"sfen_acc={v_avg['sfen/full_acc']:.3f}" ) - log.update({f"val/{k}": v for k, v in v_avg.items()}) + wandb_log.update({f"val/{k}": v for k, v in v_avg.items()}) if wandb_run is not None: - wandb_run.log(log) - - ckpt_payload = { - "epoch": epoch, - "model": model.state_dict(), - "optimizer": optimizer.state_dict(), - "backbone": args.backbone, - "hand_mode": args.hand_mode, - "image_size": args.image_size, - "hand_weight": args.hand_weight, - "wandb_run_id": wandb_run.id if wandb_run is not None else None, - } - torch.save(ckpt_payload, args.ckpt_dir / "latest.pt") - if epoch % args.save_every == 0 or epoch == args.epochs: - torch.save(ckpt_payload, args.ckpt_dir / f"epoch-{epoch:03d}.pt") + wandb_run.log(wandb_log) + + # Checkpoint save is main-only to avoid concurrent writes clobbering + # each other; the state_dict is identical across ranks after all_reduce + # in the backward pass (DDP guarantee). + if is_main(): + ckpt_payload = { + "epoch": epoch, + # Unwrap DDP and torch.compile so ckpt keys stay stable across + # launch modes (single-GPU or torchrun, --compile or not). + "model": _unwrap(model).state_dict(), + "optimizer": optimizer.state_dict(), + "scaler": scaler.state_dict() if scaler.is_enabled() else None, + "backbone": args.backbone, + "image_size": args.image_size, + "hand_weight": args.hand_weight, + "wandb_run_id": wandb_run.id if wandb_run is not None else None, + } + torch.save(ckpt_payload, args.ckpt_dir / "latest.pt") + if epoch % args.save_every == 0 or epoch == args.epochs: + torch.save(ckpt_payload, args.ckpt_dir / f"epoch-{epoch:03d}.pt") if wandb_run is not None: wandb_run.finish() # Smoke check - if args.mode == "smoke": + if args.mode == "smoke" and is_main(): final_cell = avg["board/cell_acc"] if final_cell < 0.6: print(f"[board_ocr] WARNING: did not overfit (cell_acc={final_cell:.3f}).") else: print(f"[board_ocr] smoke OK (cell_acc={final_cell:.3f}). Design pipeline verified.") + teardown_distributed() + def main() -> None: p = argparse.ArgumentParser() p.add_argument("--mode", choices=["smoke", "full"], default="smoke") p.add_argument("--backbone", default="mobilenet_v3_small", - choices=["mobilenet_v3_small", "convnext_tiny"]) + choices=[ + "mobilenet_v3_small", + "mobilenet_v3_large", + "efficientnet_b0", + "efficientnet_b1", + "convnext_atto", + "convnext_femto", + "convnext_pico", + "convnext_nano", + "convnext_tiny", + ]) p.add_argument("--pretrained", action="store_true", default=True) p.add_argument("--no-pretrained", dest="pretrained", action="store_false") - p.add_argument("--train-manifest", type=Path, default=Path("./data/ocr/train.jsonl")) - p.add_argument("--val-manifest", type=Path, default=Path("./data/ocr/val.jsonl")) - p.add_argument("--image-root", type=Path, default=Path("./data/ocr/iPhone10,1")) + p.add_argument("--train-manifest", type=Path, default=Path("./data/train.jsonl")) + p.add_argument("--val-manifest", type=Path, default=Path("./data/val.jsonl")) + p.add_argument("--image-root", type=Path, default=Path("./data/captures/d0")) p.add_argument("--hf-repo-id", type=str, default=None, help="Load from a HF dataset repo (e.g. 'ultemica/piyoshogi') instead of local manifests.") - p.add_argument("--hf-config", type=str, default="ocr_paired", - help="HF dataset config name. Use 'ocr_paired' (default, 4 images per SFEN " - "flattened to per-device examples) or a vertical config for the old schema.") - p.add_argument("--hf-devices", nargs="+", default=None, - help="Restrict paired dataset to specific device identifiers " - "(e.g. --hf-devices iPhone10,1 iPhone15,4). Default: all 4 devices.") - p.add_argument("--extra-train-source", action="append", default=[], - help="'manifest.jsonl:image_dir' pair; repeat for multi-device. " - "Concatenated with the primary --train-manifest source.") - p.add_argument("--extra-val-source", action="append", default=[], - help="'manifest.jsonl:image_dir' pair; repeat for multi-device.") p.add_argument("--image-size", type=int, default=224) p.add_argument("--limit", type=int, default=100) p.add_argument("--epochs", type=int, default=10) p.add_argument("--batch-size", type=int, default=8) p.add_argument("--num-workers", type=int, default=0) + p.add_argument("--prefetch-factor", type=int, default=4, + help="DataLoader prefetch queue depth per worker.") + p.add_argument("--preload", action="store_true", default=False, + help="Decode + resize every image once at init and hold in RAM. " + "Removes WebP decode + resize from the hot loop entirely. " + "Costs ~image_size^2*3*N*num_devices bytes " + "(~10 GB for 18k rows x 4 devices at 224).") + p.add_argument("--compile", action="store_true", default=False, + help="torch.compile the model. Worth measuring for full runs; " + "compile overhead usually not worth it for smoke runs.") p.add_argument("--lr", type=float, default=3e-4) + p.add_argument("--hand-mode", choices=["classification", "regression"], + default="classification", + help="classification: 14x19 logits + CE. regression: 14 scalars + SmoothL1.") p.add_argument("--hand-weight", type=float, default=1.0) p.add_argument("--hand-class-weight", action="store_true", default=True, - help="Use sqrt(1/freq) class weights on hand CE.") + help="Use sqrt(1/freq) class weights on hand CE (ignored under --hand-mode=regression).") p.add_argument("--no-hand-class-weight", dest="hand_class_weight", action="store_false") p.add_argument("--class-weight-clip-min", type=float, default=0.5) p.add_argument("--class-weight-clip-max", type=float, default=10.0) p.add_argument("--val-every", type=int, default=2) + p.add_argument("--log-every", type=int, default=20, + help="Log per-step train metrics to W&B every N batches.") p.add_argument("--ckpt-dir", type=Path, default=Path("./runs/board-ocr")) p.add_argument("--save-every", type=int, default=5, help="Interval for saving epoch-{N}.pt snapshots. latest.pt is saved every epoch.") diff --git a/mito_train/training/train_detector.py b/mito_train/training/train_detector.py index e009b1d..d95fcb4 100644 --- a/mito_train/training/train_detector.py +++ b/mito_train/training/train_detector.py @@ -156,7 +156,7 @@ def main() -> None: f"@0.9={metrics['iou@0.9']:.3f}" ) if run is not None: - run.log({"train/loss": train_loss, "lr": sched.get_last_lr()[0], **{f"val/{k}": v for k, v in metrics.items()}}, step=epoch) + run.log({"train/loss": train_loss, "train/lr": sched.get_last_lr()[0], **{f"val/{k}": v for k, v in metrics.items()}}, step=epoch) ck_path = args.out_dir / "latest.pt" torch.save( diff --git a/mito_train/training/train_piece.py b/mito_train/training/train_piece.py index 0a4578b..911b997 100644 --- a/mito_train/training/train_piece.py +++ b/mito_train/training/train_piece.py @@ -70,7 +70,9 @@ def _init_wandb( print("[wandb] injected CF Access headers.") settings = wandb.Settings(**settings_kwargs) if settings_kwargs else None - init_kwargs = dict(project=project, name=run_name, config=config, settings=settings) + log_dir = Path("logs") + log_dir.mkdir(parents=True, exist_ok=True) + init_kwargs = dict(project=project, name=run_name, config=config, settings=settings, dir=str(log_dir)) if run_id is not None: init_kwargs["id"] = run_id init_kwargs["resume"] = resume or "allow" @@ -172,7 +174,7 @@ def run_smoke(args: argparse.Namespace) -> None: if epoch % args.log_every == 0 or epoch in (start_epoch, args.epochs): print(f"[smoke] epoch={epoch:3d} loss={avg_loss:.4f} acc={acc:.4f}") if run is not None: - run.log({"train/loss": avg_loss, "train/acc": acc, "epoch": epoch}) + run.log({"train/loss": avg_loss, "train/acc": acc, "train/epoch": epoch}) ckpt_payload = { "epoch": epoch, diff --git a/pyproject.toml b/pyproject.toml index fc5cb3e..ef2cf36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,13 @@ [project] name = "mito-train" -version = "0.1.1" +version = "0.2.0" description = "ぴよ将棋OCR 3モデル (board-detector / piece-classifier / hand-classifier) の学習・ONNX出力" requires-python = ">=3.11" dependencies = [ "albumentations>=2.0.8", + "datasets>=5.0.0", + "hf-transfer>=0.1.9", + "huggingface-hub>=1.22.0", "matplotlib>=3.11.0", "numpy>=2.4.6", "onnx>=1.22.0", @@ -14,6 +17,7 @@ dependencies = [ "pillow>=12.3.0", "python-shogi>=1.1.1", "pyyaml>=6.0.3", + "timm>=1.0.28", "torch>=2.12.1", "torchvision>=0.27.1", "tqdm>=4.68.4", @@ -35,10 +39,3 @@ packages = ["mito_train"] # キャッシュと .venv が別ファイルシステム (マウントボリューム) なので # ハードリンク不可 → copy に固定して警告を抑止。 link-mode = "copy" - -[dependency-groups] -upload = [ - "datasets>=5.0.0", - "hf-transfer>=0.1.9", - "huggingface-hub>=1.22.0", -] diff --git a/scripts/autotune.py b/scripts/autotune.py new file mode 100644 index 0000000..a049a05 --- /dev/null +++ b/scripts/autotune.py @@ -0,0 +1,120 @@ +"""Autotune training hyperparams from detected hardware. + +Prints shell-eval-able env var assignments to stdout — the calling script +does `eval "$(python scripts/autotune.py ...)"` to import them. + +Emitted vars: + NUM_WORKERS DataLoader workers per training run + BATCH_SIZE per-GPU batch size + PREFETCH_FACTOR DataLoader prefetch depth per worker + NGPUS_DETECTED number of CUDA devices visible + VRAM_GB_PER_GPU VRAM on the first GPU (all assumed same) + CPU_COUNT host CPU count + +Sizing rules (deliberately conservative — tuner should adjust up if VRAM +util stays low mid-run): + - batch = ~ vram_gb * scale / footprint(backbone, image_size) + where footprint approximates activation + optimizer memory per sample + - workers = min(cpu_count / concurrent_runs, 16) + with concurrent_runs being 1 in sequential/DDP modes, NGPUS in parallel mode +""" +from __future__ import annotations +import argparse +import multiprocessing +import sys + + +# Per-sample VRAM footprint at 224x224 batch=1 (GB), rough empirical estimate. +# Scales quadratically with image_size (H*W). +BACKBONE_FOOTPRINT_GB = { + "mobilenet_v3_small": 0.020, + "mobilenet_v3_large": 0.035, + "efficientnet_b0": 0.045, + "efficientnet_b1": 0.055, + "convnext_atto": 0.040, + "convnext_femto": 0.050, + "convnext_pico": 0.070, + "convnext_nano": 0.090, + "convnext_tiny": 0.130, +} + + +def detect_gpu() -> tuple[int, float]: + try: + import torch + except ImportError: + return 0, 0.0 + if not torch.cuda.is_available(): + return 0, 0.0 + n = torch.cuda.device_count() + if n == 0: + return 0, 0.0 + props = torch.cuda.get_device_properties(0) + vram_gb = props.total_memory / (1024 ** 3) + return n, vram_gb + + +def suggest_batch(backbone: str, image_size: int, vram_gb: float, + concurrent_on_gpu: int = 1) -> int: + """Round to a nearby power of 2 for kernel friendliness.""" + if vram_gb <= 0: + return 32 # CPU fallback + footprint = BACKBONE_FOOTPRINT_GB.get(backbone, 0.10) + footprint *= (image_size / 224) ** 2 + # ~60% of VRAM budgeted for activations+batch (rest: model, optimizer, preload) + usable_vram = vram_gb * 0.60 / concurrent_on_gpu + raw = usable_vram / footprint + # Snap to nearest power of two below raw, clamped to a sane range. + b = 32 + while b * 2 <= raw and b < 1024: + b *= 2 + return max(16, b) + + +def suggest_workers(cpu_count: int, concurrent_runs: int) -> int: + """Leave a couple cores free for the main + DataLoader main; cap at 16.""" + per_run = max(1, (cpu_count - 2) // max(1, concurrent_runs)) + return min(per_run, 16) + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument("--backbone", required=True, + help="Backbone name; used to look up per-sample footprint.") + p.add_argument("--image-size", type=int, default=224) + p.add_argument("--mode", choices=["sequential", "ddp", "parallel"], + default="sequential", + help="Launch mode. Affects workers division and concurrent-per-GPU count.") + p.add_argument("--nprocs", type=int, default=None, + help="For DDP: GPUs used per run. For parallel: GPUs sharing the sweep. " + "Default = auto-detected GPU count.") + args = p.parse_args() + + ngpus, vram_gb = detect_gpu() + cpu_count = multiprocessing.cpu_count() + + if args.mode == "parallel": + concurrent_runs = args.nprocs or ngpus or 1 + concurrent_on_gpu = 1 # one backbone per GPU + elif args.mode == "ddp": + concurrent_runs = 1 + concurrent_on_gpu = 1 # single backbone across many GPUs + else: + concurrent_runs = 1 + concurrent_on_gpu = 1 + + batch = suggest_batch(args.backbone, args.image_size, vram_gb, concurrent_on_gpu) + workers = suggest_workers(cpu_count, concurrent_runs) + prefetch = 4 if workers <= 4 else 8 + + print(f"NUM_WORKERS={workers}") + print(f"BATCH_SIZE={batch}") + print(f"PREFETCH_FACTOR={prefetch}") + print(f"NGPUS_DETECTED={ngpus}") + print(f"VRAM_GB_PER_GPU={vram_gb:.1f}") + print(f"CPU_COUNT={cpu_count}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/train_backbones.sh b/scripts/train_backbones.sh new file mode 100755 index 0000000..7c753ad --- /dev/null +++ b/scripts/train_backbones.sh @@ -0,0 +1,298 @@ +#!/usr/bin/env bash +# Backbone sweep: train each backbone for the same number of epochs and settings, +# so W&B side-by-side lets you spot the params-vs-accuracy Pareto front. +# +# Three launch modes; pick one via env vars: +# (default) 1 run at a time, single GPU +# NPROC_PER_NODE=N 1 run at a time, N-GPU DDP (torchrun) +# PARALLEL_GPU=1 many runs concurrently, 1 GPU each (best for a +# multi-backbone sweep on a multi-GPU node) +# +# General env overrides: +# EPOCHS BATCH_SIZE NUM_WORKERS PREFETCH_FACTOR IMAGE_SIZE LR +# HF_REPO_ID (defaults to ultemica/piyoshogi) +# CKPT_ROOT (dir under which each backbone gets its own .//) +# BACKBONES (space-separated subset; defaults to all 9, small->large) +# SKIP_EXISTING=1 skip a backbone if ckpt_dir/latest.pt exists (blunt) +# RESUME_INCOMPLETE=1 smart: skip if epoch-EPOCHS.pt exists (done), resume +# if only latest.pt exists (partial), else fresh +# COMPILE=1 pass --compile to each run +# +# Parallel-mode extras: +# NGPUS= override auto-detected GPU count +# PARALLEL_LOG_DIR directory for per-backbone stdout logs +# (defaults to CKPT_ROOT/logs) +# +# Examples: +# ./scripts/train_backbones.sh +# EPOCHS=30 BACKBONES="mobilenet_v3_small convnext_atto" ./scripts/train_backbones.sh +# RESUME_INCOMPLETE=1 ./scripts/train_backbones.sh # resume mid-sweep +# NPROC_PER_NODE=8 ./scripts/train_backbones.sh # DDP each backbone across 8 GPUs +# PARALLEL_GPU=1 ./scripts/train_backbones.sh # 8 backbones at once, 1 GPU each + +set -euo pipefail + +EPOCHS=${EPOCHS:-50} +BATCH_SIZE=${BATCH_SIZE:-32} +NUM_WORKERS=${NUM_WORKERS:-16} +PREFETCH_FACTOR=${PREFETCH_FACTOR:-8} +IMAGE_SIZE=${IMAGE_SIZE:-224} +LR=${LR:-3e-4} +HF_REPO_ID=${HF_REPO_ID:-ultemica/piyoshogi} +CKPT_ROOT=${CKPT_ROOT:-./runs} +SKIP_EXISTING=${SKIP_EXISTING:-0} +RESUME_INCOMPLETE=${RESUME_INCOMPLETE:-0} +COMPILE=${COMPILE:-0} +NPROC_PER_NODE=${NPROC_PER_NODE:-1} +MASTER_PORT=${MASTER_PORT:-29500} +PARALLEL_GPU=${PARALLEL_GPU:-0} +PARALLEL_LOG_DIR=${PARALLEL_LOG_DIR:-${CKPT_ROOT}/logs} + +# Ordered small -> large so early results inform whether we need the heavies. +DEFAULT_BACKBONES=( + mobilenet_v3_small + mobilenet_v3_large + convnext_atto + convnext_femto + efficientnet_b0 + convnext_pico + efficientnet_b1 + convnext_nano + convnext_tiny +) +if [[ -n "${BACKBONES:-}" ]]; then + read -r -a BACKBONE_LIST <<< "${BACKBONES}" +else + BACKBONE_LIST=("${DEFAULT_BACKBONES[@]}") +fi + +extra_args=(--preload) +if [[ "${COMPILE}" == "1" ]]; then + extra_args+=(--compile) +fi + +# Mode-guard: parallel and DDP are mutually exclusive (parallel implies 1 GPU per run). +if [[ "${PARALLEL_GPU}" == "1" && "${NPROC_PER_NODE}" -gt 1 ]]; then + echo "[sweep] error: PARALLEL_GPU=1 and NPROC_PER_NODE>1 are mutually exclusive." >&2 + exit 2 +fi + +detect_gpu_count() { + uv run python -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1 +} +NGPUS=${NGPUS:-$(detect_gpu_count)} + +echo "[sweep] epochs=${EPOCHS} batch=${BATCH_SIZE} workers=${NUM_WORKERS} image_size=${IMAGE_SIZE} lr=${LR}" +echo "[sweep] backbones: ${BACKBONE_LIST[*]}" +echo "[sweep] ckpt root: ${CKPT_ROOT}" +if [[ "${PARALLEL_GPU}" == "1" ]]; then + echo "[sweep] parallel-GPU: ${NGPUS} GPUs, one backbone per GPU" + echo "[sweep] per-run logs: ${PARALLEL_LOG_DIR}" +elif [[ "${NPROC_PER_NODE}" -gt 1 ]]; then + echo "[sweep] multi-GPU DDP: nproc_per_node=${NPROC_PER_NODE} (torchrun)" + echo "[sweep] effective batch = ${BATCH_SIZE} x ${NPROC_PER_NODE} = $(( BATCH_SIZE * NPROC_PER_NODE ))" +fi +echo + +declare -a completed=() +declare -a skipped=() +declare -a failed=() + +print_summary() { + echo + echo "[sweep] === summary ===" + echo "[sweep] completed: ${completed[*]:-}" + echo "[sweep] skipped: ${skipped[*]:-}" + echo "[sweep] failed: ${failed[*]:-}" +} + +# Ctrl-C mid-loop should stop cleanly but still print a summary. +# In parallel mode, also kill any running child processes. +declare -a child_pids=() +cleanup_children() { + for pid in "${child_pids[@]}"; do + kill "${pid}" 2>/dev/null || true + done +} +trap 'echo; echo "[sweep] interrupted"; cleanup_children; print_summary; exit 130' INT + +final_ckpt_name=$(printf "epoch-%03d.pt" "${EPOCHS}") + +# resolve_run_state — echoes one of: +# SKIP (final ckpt exists in RESUME_INCOMPLETE; or latest.pt in SKIP_EXISTING) +# RESUME (partial ckpt found in RESUME_INCOMPLETE) +# FRESH (start from epoch 1) +resolve_run_state() { + local latest="$1" final="$2" + if [[ "${RESUME_INCOMPLETE}" == "1" ]]; then + if [[ -f "${final}" ]]; then echo SKIP; return; fi + if [[ -f "${latest}" ]]; then echo RESUME; return; fi + echo FRESH; return + fi + if [[ "${SKIP_EXISTING}" == "1" && -f "${latest}" ]]; then + echo SKIP; return + fi + echo FRESH +} + +# run_single BACKBONE — launches a run in the FOREGROUND. Used by sequential +# and DDP modes. Returns the training exit code. +run_single() { + local BB="$1" + local ckpt_dir="${CKPT_ROOT}/board-ocr-${BB}" + local latest_ckpt="${ckpt_dir}/latest.pt" + local final_ckpt="${ckpt_dir}/${final_ckpt_name}" + + local state + state=$(resolve_run_state "${latest_ckpt}" "${final_ckpt}") + + local per_run_extra=() + case "${state}" in + SKIP) + echo "[sweep] SKIP ${BB}" + skipped+=("${BB}") + return 0 + ;; + RESUME) + echo "[sweep] RESUME ${BB} from ${latest_ckpt}" + per_run_extra+=(--resume latest) + ;; + esac + + echo + echo "[sweep] ============================================================" + echo "[sweep] ${BB} -> ${ckpt_dir}" + echo "[sweep] ============================================================" + local start_ts + start_ts=$(date +%s) + + local launcher=() + if [[ "${NPROC_PER_NODE}" -gt 1 ]]; then + launcher=(uv run torchrun + --standalone + --nproc_per_node="${NPROC_PER_NODE}" + --master_port="${MASTER_PORT}" + -m mito_train.training.train_board_ocr) + else + launcher=(uv run python -m mito_train.training.train_board_ocr) + fi + + if "${launcher[@]}" \ + --hf-repo-id "${HF_REPO_ID}" \ + --mode full \ + --backbone "${BB}" \ + --image-size "${IMAGE_SIZE}" \ + --epochs "${EPOCHS}" \ + --batch-size "${BATCH_SIZE}" \ + --num-workers "${NUM_WORKERS}" \ + --prefetch-factor "${PREFETCH_FACTOR}" \ + --lr "${LR}" \ + --ckpt-dir "${ckpt_dir}" \ + "${extra_args[@]}" \ + "${per_run_extra[@]}" + then + local elapsed=$(( $(date +%s) - start_ts )) + echo "[sweep] ${BB} DONE in ${elapsed}s" + completed+=("${BB}") + else + local rc=$? + echo "[sweep] ${BB} FAILED (exit ${rc})" + failed+=("${BB}") + fi +} + +# launch_parallel BACKBONE GPU_IDX — starts a background run pinned to one GPU. +# Its stdout+stderr is redirected to per-backbone log file. Sets $!. +launch_parallel() { + local BB="$1" gpu="$2" + local ckpt_dir="${CKPT_ROOT}/board-ocr-${BB}" + local latest_ckpt="${ckpt_dir}/latest.pt" + local final_ckpt="${ckpt_dir}/${final_ckpt_name}" + local log_file="${PARALLEL_LOG_DIR}/${BB}.log" + + local state + state=$(resolve_run_state "${latest_ckpt}" "${final_ckpt}") + + local per_run_extra=() + case "${state}" in + SKIP) + echo "[sweep] SKIP ${BB}" + skipped+=("${BB}") + return 1 # signal caller: no process started + ;; + RESUME) + echo "[sweep] RESUME ${BB} on GPU ${gpu} (log: ${log_file})" + per_run_extra+=(--resume latest) + ;; + FRESH) + echo "[sweep] START ${BB} on GPU ${gpu} (log: ${log_file})" + ;; + esac + + # CUDA_VISIBLE_DEVICES pins this subprocess to a single GPU without touching + # the parent env. torch inside the child sees exactly one device as cuda:0. + CUDA_VISIBLE_DEVICES="${gpu}" \ + uv run python -m mito_train.training.train_board_ocr \ + --hf-repo-id "${HF_REPO_ID}" \ + --mode full \ + --backbone "${BB}" \ + --image-size "${IMAGE_SIZE}" \ + --epochs "${EPOCHS}" \ + --batch-size "${BATCH_SIZE}" \ + --num-workers "${NUM_WORKERS}" \ + --prefetch-factor "${PREFETCH_FACTOR}" \ + --lr "${LR}" \ + --ckpt-dir "${ckpt_dir}" \ + "${extra_args[@]}" \ + "${per_run_extra[@]}" \ + >"${log_file}" 2>&1 & + return 0 +} + +if [[ "${PARALLEL_GPU}" == "1" ]]; then + mkdir -p "${PARALLEL_LOG_DIR}" + + # Wave-based scheduler: launch up to NGPUS runs, wait for all to finish, + # then start the next wave. Simple and predictable; a continuous + # dispatcher would be slightly faster but hides failure modes. + idx=0 + total=${#BACKBONE_LIST[@]} + wave=0 + while (( idx < total )); do + wave=$(( wave + 1 )) + echo + echo "[sweep] --- wave ${wave} ---" + wave_pids=() + wave_backbones=() + for (( g=0; g/dev/null || echo 1 +} + +NPROC_PER_NODE=${NPROC_PER_NODE:-$(detect_gpu_count)} +MASTER_PORT=${MASTER_PORT:-29500} +PER_GPU_WORKERS=${PER_GPU_WORKERS:-4} + +EPOCHS=${EPOCHS:-50} +BATCH_SIZE=${BATCH_SIZE:-32} +PREFETCH_FACTOR=${PREFETCH_FACTOR:-8} +IMAGE_SIZE=${IMAGE_SIZE:-224} +LR=${LR:-3e-4} +BACKBONE=${BACKBONE:-mobilenet_v3_small} +HF_REPO_ID=${HF_REPO_ID:-ultemica/piyoshogi} +CKPT_DIR=${CKPT_DIR:-./runs/board-ocr-${BACKBONE}} +SAVE_EVERY=${SAVE_EVERY:-5} + +extra_args=(--preload) +if [[ -n "${RESUME:-}" ]]; then + extra_args+=(--resume "${RESUME}") +fi +if [[ -n "${WANDB_RUN_ID:-}" ]]; then + extra_args+=(--wandb-run-id "${WANDB_RUN_ID}") +fi +if [[ "${COMPILE:-0}" == "1" ]]; then + extra_args+=(--compile) +fi + +echo "[train_multi_gpu] nproc=${NPROC_PER_NODE} backbone=${BACKBONE} batch=${BATCH_SIZE} (per rank)" +echo "[train_multi_gpu] effective batch = ${BATCH_SIZE} x ${NPROC_PER_NODE} = $(( BATCH_SIZE * NPROC_PER_NODE ))" +echo "[train_multi_gpu] workers=${PER_GPU_WORKERS} per rank ($(( PER_GPU_WORKERS * NPROC_PER_NODE )) total)" +echo "[train_multi_gpu] ckpt_dir=${CKPT_DIR}" + +# --standalone: single-node rendezvous, avoids setting MASTER_ADDR/MASTER_PORT +# manually for the common single-host case. +uv run torchrun \ + --standalone \ + --nproc_per_node="${NPROC_PER_NODE}" \ + --master_port="${MASTER_PORT}" \ + -m mito_train.training.train_board_ocr \ + --hf-repo-id "${HF_REPO_ID}" \ + --mode full \ + --backbone "${BACKBONE}" \ + --image-size "${IMAGE_SIZE}" \ + --epochs "${EPOCHS}" \ + --batch-size "${BATCH_SIZE}" \ + --num-workers "${PER_GPU_WORKERS}" \ + --prefetch-factor "${PREFETCH_FACTOR}" \ + --lr "${LR}" \ + --ckpt-dir "${CKPT_DIR}" \ + --save-every "${SAVE_EVERY}" \ + "${extra_args[@]}" diff --git a/uv.lock b/uv.lock index 4d6f254..f017df5 100644 --- a/uv.lock +++ b/uv.lock @@ -1119,10 +1119,13 @@ wheels = [ [[package]] name = "mito-train" -version = "0.1.1" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "albumentations" }, + { name = "datasets" }, + { name = "hf-transfer" }, + { name = "huggingface-hub" }, { name = "matplotlib" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, @@ -1133,6 +1136,7 @@ dependencies = [ { name = "pillow" }, { name = "python-shogi" }, { name = "pyyaml" }, + { name = "timm" }, { name = "torch" }, { name = "torchvision" }, { name = "tqdm" }, @@ -1143,16 +1147,12 @@ experiment = [ { name = "wandb" }, ] -[package.dev-dependencies] -upload = [ - { name = "datasets" }, - { name = "hf-transfer" }, - { name = "huggingface-hub" }, -] - [package.metadata] requires-dist = [ { name = "albumentations", specifier = ">=2.0.8" }, + { name = "datasets", specifier = ">=5.0.0" }, + { name = "hf-transfer", specifier = ">=0.1.9" }, + { name = "huggingface-hub", specifier = ">=1.22.0" }, { name = "matplotlib", specifier = ">=3.11.0" }, { name = "numpy", specifier = ">=2.4.6" }, { name = "onnx", specifier = ">=1.22.0" }, @@ -1162,6 +1162,7 @@ requires-dist = [ { name = "pillow", specifier = ">=12.3.0" }, { name = "python-shogi", specifier = ">=1.1.1" }, { name = "pyyaml", specifier = ">=6.0.3" }, + { name = "timm", specifier = ">=1.0.28" }, { name = "torch", specifier = ">=2.12.1" }, { name = "torchvision", specifier = ">=0.27.1" }, { name = "tqdm", specifier = ">=4.68.4" }, @@ -1169,13 +1170,6 @@ requires-dist = [ ] provides-extras = ["experiment"] -[package.metadata.requires-dev] -upload = [ - { name = "datasets", specifier = ">=5.0.0" }, - { name = "hf-transfer", specifier = ">=0.1.9" }, - { name = "huggingface-hub", specifier = ">=1.22.0" }, -] - [[package]] name = "ml-dtypes" version = "0.5.4" @@ -2354,6 +2348,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, ] +[[package]] +name = "safetensors" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, +] + [[package]] name = "scipy" version = "1.17.1" @@ -2676,6 +2694,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "timm" +version = "1.0.28" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "huggingface-hub" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch" }, + { name = "torchvision" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/03/e41389ac641747bfec48d016fde8be1eade1901e6f2c1aedcb0c8cb4b5d9/timm-1.0.28.tar.gz", hash = "sha256:3789d313fdd5541a327b60180d70dbb4bdec73db8ff0655e413db3c3d134a9a4", size = 2451413, upload-time = "2026-07-11T17:24:32.615Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/76/de1bfac17d183c49c6d0887903d3064ced51cf1d9ba7a8d611c1a8808c4f/timm-1.0.28-py3-none-any.whl", hash = "sha256:e577b88da96b3a722ea5e2f042455ce6f715d398304d8e63b17d126ed7d89968", size = 2597944, upload-time = "2026-07-11T17:24:30.869Z" }, +] + [[package]] name = "torch" version = "2.12.1"