From a68bec56678e4ee76d788874f50b06ccb68ddfb6 Mon Sep 17 00:00:00 2001 From: tkgstrator Date: Mon, 13 Jul 2026 12:19:41 +0900 Subject: [PATCH 1/6] chore(train): scope wandb project by image size and package version (#8) wandb project name for BoardOCR was a fixed "mito-train-board-ocr", which collapsed runs from different image_size and package versions into a single board on wandb. Interpolate image_size and mito-train pyproject version so each combination gets its own project (e.g. mito-train-board-ocr-w224-v0.2.0). Co-authored-by: Claude Opus 4.8 (1M context) --- docs/backbones.md | 2 +- docs/ocr-scaling-outlook.md | 2 +- mito_train/training/train_board_ocr.py | 4 +++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/backbones.md b/docs/backbones.md index ecab32a..35494db 100644 --- a/docs/backbones.md +++ b/docs/backbones.md @@ -364,5 +364,5 @@ BoardOCR では未実装だが、モバイル配信で精度を詰めたい場 各 backbone の val cell_acc / sfen_full_acc / wall-clock は sweep 完了後に W&B で並び、ここに引用予定。 -- W&B project: `mito-train-board-ocr` +- W&B project: `mito-train-board-ocr-w{image_size}-v{package_version}` (例: `mito-train-board-ocr-w224-v0.2.0`) - run 名は各 backbone 名そのまま（例: `mobilenet_v3_small`, `convnext_tiny`） diff --git a/docs/ocr-scaling-outlook.md b/docs/ocr-scaling-outlook.md index a4fc19f..3eb6844 100644 --- a/docs/ocr-scaling-outlook.md +++ b/docs/ocr-scaling-outlook.md @@ -257,7 +257,7 @@ EPOCHS=50 IMAGE_SIZE=288 BACKBONES="mobilenet_v3_small convnext_atto convnext_ti RESUME_INCOMPLETE=1 EPOCHS=100 IMAGE_SIZE=288 ./scripts/train_backbones.sh ``` -W&B project: `mito-train-board-ocr`。今回の sweep run 一覧： +W&B project: `mito-train-board-ocr`（当時。以降は `mito-train-board-ocr-w{image_size}-v{package_version}` に分割）。今回の sweep run 一覧： | Backbone | Run ID | |---|---| diff --git a/mito_train/training/train_board_ocr.py b/mito_train/training/train_board_ocr.py index 6287f47..e2dc504 100644 --- a/mito_train/training/train_board_ocr.py +++ b/mito_train/training/train_board_ocr.py @@ -17,6 +17,7 @@ from __future__ import annotations import argparse +from importlib.metadata import version as _pkg_version from pathlib import Path import torch @@ -301,8 +302,9 @@ def log_main(msg: str) -> None: # Only rank 0 talks to W&B; other ranks keep wandb_run=None and log nothing. run_name = args.backbone + wandb_project = f"mito-train-board-ocr-w{args.image_size}-v{_pkg_version('mito-train')}" wandb_run = _init_wandb( - project="mito-train-board-ocr", + project=wandb_project, run_name=run_name, run_id=resumed_wandb_id, resume="allow" if resumed_wandb_id else None, From 577a5103ebbff7d07096b0f245efad92fbc8bb4a Mon Sep 17 00:00:00 2001 From: tkgstrator Date: Mon, 13 Jul 2026 12:41:57 +0900 Subject: [PATCH 2/6] feat(train_board_ocr): replace --ckpt-dir with --model-name and add hand MAE (#9) - Replace --ckpt-dir with --model-name (+ --ckpt-root, default ./runs). Checkpoints now land at //. model_name defaults to board-ocr-, matching the sweep naming. - --resume latest no longer errors when latest.pt is missing; it falls through to a fresh run so sweep scripts can pass it unconditionally. - Add hand/mae to metrics and surface board_acc, hand_acc, and hand_mae in train + val log lines. MAE treats the count head as ordinal, so "off by 1" reads very differently from "off by 9". - Update train.sh, train_multi_gpu.sh, train_backbones.sh to pass --model-name. - Bump version to 0.3.0 (CLI flag rename is breaking; pre-1.0 => minor). Co-authored-by: Claude Opus 4.8 (1M context) --- mito_train/training/train_board_ocr.py | 60 ++++++++++++++++++-------- pyproject.toml | 2 +- scripts/train.sh | 6 +-- scripts/train_backbones.sh | 10 +++-- scripts/train_multi_gpu.sh | 6 +-- uv.lock | 2 +- 6 files changed, 55 insertions(+), 31 deletions(-) diff --git a/mito_train/training/train_board_ocr.py b/mito_train/training/train_board_ocr.py index e2dc504..fab556a 100644 --- a/mito_train/training/train_board_ocr.py +++ b/mito_train/training/train_board_ocr.py @@ -102,7 +102,7 @@ def compute_loss( METRIC_KEYS = ( "board/cell_acc", "board/full_acc", - "hand/slot_acc", "hand/full_acc", "sfen/full_acc", + "hand/slot_acc", "hand/full_acc", "hand/mae", "sfen/full_acc", ) @@ -131,6 +131,9 @@ def compute_metrics( hand_slot_acc = hand_ok.float().mean() hand_all = hand_ok.all(dim=1) hand_full = hand_all.float().mean() + # MAE over all 14 slots — treats the count head as ordinal so + # "off by 1" is much better than "off by 9". + hand_mae = (hand_pred - hand_target).abs().float().mean() sfen_full = (board_all & hand_all).float().mean() return { @@ -138,6 +141,7 @@ def compute_metrics( "board/full_acc": board_full, "hand/slot_acc": hand_slot_acc, "hand/full_acc": hand_full, + "hand/mae": hand_mae, "sfen/full_acc": sfen_full, } @@ -276,25 +280,31 @@ def log_main(msg: str) -> None: model.parameters(), lr=args.lr, weight_decay=1e-4, fused=use_cuda, ) + ckpt_dir = args.ckpt_root / args.model_name if is_main(): - args.ckpt_dir.mkdir(parents=True, exist_ok=True) + 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 - log_main(f"[board_ocr] resume from {ckpt_path}") - ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) - _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") - log_main(f"[board_ocr] resumed at epoch={start_epoch} (ckpt was epoch {ckpt['epoch']})") - if resumed_wandb_id: - log_main(f"[board_ocr] will resume wandb run id={resumed_wandb_id}") + ckpt_path = ckpt_dir / "latest.pt" if str(args.resume) == "latest" else args.resume + # `--resume latest` on a fresh model_name has no checkpoint to load — treat + # as fresh start so sweep scripts can pass --resume latest unconditionally. + if str(args.resume) == "latest" and not ckpt_path.exists(): + log_main(f"[board_ocr] --resume latest requested but {ckpt_path} not found — starting fresh") + else: + log_main(f"[board_ocr] resume from {ckpt_path}") + ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) + _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") + log_main(f"[board_ocr] resumed at epoch={start_epoch} (ckpt was epoch {ckpt['epoch']})") + if 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 @@ -408,6 +418,9 @@ def log_main(msg: str) -> None: 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"board_acc={avg['board/full_acc']:.3f} " + f"hand_acc={avg['hand/full_acc']:.3f} " + f"hand_mae={avg['hand/mae']:.3f} " f"sfen_acc={avg['sfen/full_acc']:.3f}" ) @@ -448,6 +461,9 @@ def log_main(msg: str) -> None: } log_main( f"[board_ocr] val cell_acc={v_avg['board/cell_acc']:.3f} " + f"board_acc={v_avg['board/full_acc']:.3f} " + f"hand_acc={v_avg['hand/full_acc']:.3f} " + f"hand_mae={v_avg['hand/mae']:.3f} " f"sfen_acc={v_avg['sfen/full_acc']:.3f}" ) wandb_log.update({f"val/{k}": v for k, v in v_avg.items()}) @@ -471,9 +487,9 @@ def log_main(msg: str) -> None: "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") + torch.save(ckpt_payload, 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") + torch.save(ckpt_payload, ckpt_dir / f"epoch-{epoch:03d}.pt") if wandb_run is not None: wandb_run.finish() @@ -539,14 +555,20 @@ def main() -> None: 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("--model-name", type=str, default=None, + help="Run identifier. Checkpoints go to //. " + "Defaults to 'board-ocr-'.") + p.add_argument("--ckpt-root", type=Path, default=Path("./runs"), + help="Directory under which each run gets its own / subdir.") p.add_argument("--save-every", type=int, default=5, help="Interval for saving epoch-{N}.pt snapshots. latest.pt is saved every epoch.") p.add_argument("--resume", type=Path, default=None, - help="Checkpoint path. Pass 'latest' to load --ckpt-dir/latest.pt.") + help="Checkpoint path. Pass 'latest' to load ./runs//latest.pt.") p.add_argument("--wandb-run-id", type=str, default=None, help="Force resume this W&B run id (overrides ckpt's stored id).") args = p.parse_args() + if args.model_name is None: + args.model_name = f"board-ocr-{args.backbone}" run(args) diff --git a/pyproject.toml b/pyproject.toml index ef2cf36..65b9dfb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mito-train" -version = "0.2.0" +version = "0.3.0" description = "ぴよ将棋OCR 3モデル (board-detector / piece-classifier / hand-classifier) の学習・ONNX出力" requires-python = ">=3.11" dependencies = [ diff --git a/scripts/train.sh b/scripts/train.sh index a422457..64116e8 100755 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # board OCR full 学習: mobilenet_v3_small + Apple Silicon MPS 前提。 # 環境変数で上書き可: EPOCHS, BATCH_SIZE, NUM_WORKERS, LR, BACKBONE, IMAGE_SIZE, -# CKPT_DIR, SAVE_EVERY, RESUME, WANDB_RUN_ID, HF_REPO_ID +# MODEL_NAME, SAVE_EVERY, RESUME, WANDB_RUN_ID, HF_REPO_ID # 途中再開したいときは: RESUME=latest EPOCHS=30 ./scripts/train.sh # 過去の W&B run に強制的に紐づけたいときは: # WANDB_RUN_ID=xxxxxxxx RESUME=latest EPOCHS=30 ./scripts/train.sh @@ -15,7 +15,7 @@ NUM_WORKERS=${NUM_WORKERS:-8} LR=${LR:-6e-4} BACKBONE=${BACKBONE:-mobilenet_v3_small} IMAGE_SIZE=${IMAGE_SIZE:-288} -CKPT_DIR=${CKPT_DIR:-./runs/board-ocr-v2} +MODEL_NAME=${MODEL_NAME:-board-ocr-v2} SAVE_EVERY=${SAVE_EVERY:-5} extra_args=() @@ -37,6 +37,6 @@ uv run python -m mito_train.training.train_board_ocr \ --batch-size "${BATCH_SIZE}" \ --num-workers "${NUM_WORKERS}" \ --lr "${LR}" \ - --ckpt-dir "${CKPT_DIR}" \ + --model-name "${MODEL_NAME}" \ --save-every "${SAVE_EVERY}" \ "${extra_args[@]}" diff --git a/scripts/train_backbones.sh b/scripts/train_backbones.sh index 7c753ad..1aa74f7 100755 --- a/scripts/train_backbones.sh +++ b/scripts/train_backbones.sh @@ -139,7 +139,8 @@ resolve_run_state() { # and DDP modes. Returns the training exit code. run_single() { local BB="$1" - local ckpt_dir="${CKPT_ROOT}/board-ocr-${BB}" + local model_name="board-ocr-${BB}" + local ckpt_dir="${CKPT_ROOT}/${model_name}" local latest_ckpt="${ckpt_dir}/latest.pt" local final_ckpt="${ckpt_dir}/${final_ckpt_name}" @@ -187,7 +188,7 @@ run_single() { --num-workers "${NUM_WORKERS}" \ --prefetch-factor "${PREFETCH_FACTOR}" \ --lr "${LR}" \ - --ckpt-dir "${ckpt_dir}" \ + --model-name "${model_name}" \ "${extra_args[@]}" \ "${per_run_extra[@]}" then @@ -205,7 +206,8 @@ run_single() { # 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 model_name="board-ocr-${BB}" + local ckpt_dir="${CKPT_ROOT}/${model_name}" local latest_ckpt="${ckpt_dir}/latest.pt" local final_ckpt="${ckpt_dir}/${final_ckpt_name}" local log_file="${PARALLEL_LOG_DIR}/${BB}.log" @@ -242,7 +244,7 @@ launch_parallel() { --num-workers "${NUM_WORKERS}" \ --prefetch-factor "${PREFETCH_FACTOR}" \ --lr "${LR}" \ - --ckpt-dir "${ckpt_dir}" \ + --model-name "${model_name}" \ "${extra_args[@]}" \ "${per_run_extra[@]}" \ >"${log_file}" 2>&1 & diff --git a/scripts/train_multi_gpu.sh b/scripts/train_multi_gpu.sh index 41ace65..04360b7 100755 --- a/scripts/train_multi_gpu.sh +++ b/scripts/train_multi_gpu.sh @@ -31,7 +31,7 @@ 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}} +MODEL_NAME=${MODEL_NAME:-board-ocr-${BACKBONE}} SAVE_EVERY=${SAVE_EVERY:-5} extra_args=(--preload) @@ -48,7 +48,7 @@ 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}" +echo "[train_multi_gpu] model_name=${MODEL_NAME} (ckpt_dir=./runs/${MODEL_NAME})" # --standalone: single-node rendezvous, avoids setting MASTER_ADDR/MASTER_PORT # manually for the common single-host case. @@ -66,6 +66,6 @@ uv run torchrun \ --num-workers "${PER_GPU_WORKERS}" \ --prefetch-factor "${PREFETCH_FACTOR}" \ --lr "${LR}" \ - --ckpt-dir "${CKPT_DIR}" \ + --model-name "${MODEL_NAME}" \ --save-every "${SAVE_EVERY}" \ "${extra_args[@]}" diff --git a/uv.lock b/uv.lock index f017df5..348ec64 100644 --- a/uv.lock +++ b/uv.lock @@ -1119,7 +1119,7 @@ wheels = [ [[package]] name = "mito-train" -version = "0.2.0" +version = "0.3.0" source = { editable = "." } dependencies = [ { name = "albumentations" }, From 2c03c414a5da222b14249e740a5ddfbf561743d7 Mon Sep 17 00:00:00 2001 From: tkgstrator Date: Mon, 13 Jul 2026 12:52:33 +0900 Subject: [PATCH 3/6] build(train): add python-dotenv for .env credential loading (#10) Load environment variables from .env into os.environ at the start of each training entrypoint (train_board_ocr, train_piece, train_detector) so HuggingFace (HF_TOKEN) and Weights & Biases (WANDB_API_KEY, CF_ACCESS_CLIENT_*) can resolve credentials without exporting them by hand. load_dotenv() is a no-op when .env is missing and does not override existing os.environ entries, so devcontainer-exported values still win. Bump version to 0.3.1. Co-authored-by: Claude Opus 4.8 (1M context) --- mito_train/training/train_board_ocr.py | 5 +++++ mito_train/training/train_detector.py | 5 +++++ mito_train/training/train_piece.py | 5 +++++ pyproject.toml | 3 ++- uv.lock | 13 ++++++++++++- 5 files changed, 29 insertions(+), 2 deletions(-) diff --git a/mito_train/training/train_board_ocr.py b/mito_train/training/train_board_ocr.py index fab556a..07b8385 100644 --- a/mito_train/training/train_board_ocr.py +++ b/mito_train/training/train_board_ocr.py @@ -20,6 +20,8 @@ from importlib.metadata import version as _pkg_version from pathlib import Path +from dotenv import load_dotenv + import torch import torch.nn as nn import torch.nn.functional as F @@ -506,6 +508,9 @@ def log_main(msg: str) -> None: def main() -> None: + # Pull HF_TOKEN / WANDB_API_KEY / CF_* out of .env into os.environ before + # any HF or wandb call resolves credentials. No-op if .env is missing. + load_dotenv() p = argparse.ArgumentParser() p.add_argument("--mode", choices=["smoke", "full"], default="smoke") p.add_argument("--backbone", default="mobilenet_v3_small", diff --git a/mito_train/training/train_detector.py b/mito_train/training/train_detector.py index d95fcb4..cd0d953 100644 --- a/mito_train/training/train_detector.py +++ b/mito_train/training/train_detector.py @@ -20,6 +20,8 @@ import argparse from pathlib import Path +from dotenv import load_dotenv + import torch import torch.nn as nn import torch.nn.functional as F @@ -74,6 +76,9 @@ def evaluate(model: nn.Module, loader: DataLoader, device: str) -> dict[str, flo def main() -> None: + # Pull HF_TOKEN / WANDB_API_KEY / CF_* out of .env into os.environ before + # any wandb or HF call resolves credentials. No-op if .env is missing. + load_dotenv() p = argparse.ArgumentParser() p.add_argument("--train-manifest", type=Path, default=Path("data/detector/train.jsonl")) p.add_argument("--val-manifest", type=Path, default=Path("data/detector/val.jsonl")) diff --git a/mito_train/training/train_piece.py b/mito_train/training/train_piece.py index 911b997..dd6a5b7 100644 --- a/mito_train/training/train_piece.py +++ b/mito_train/training/train_piece.py @@ -18,6 +18,8 @@ import os from pathlib import Path +from dotenv import load_dotenv + import torch import torch.nn as nn from torch.utils.data import DataLoader @@ -219,6 +221,9 @@ def run_manifest(args: argparse.Namespace) -> None: def main() -> None: + # Pull WANDB_API_KEY / CF_* / HF_TOKEN out of .env into os.environ before + # any wandb or HF call resolves credentials. No-op if .env is missing. + load_dotenv() p = argparse.ArgumentParser() p.add_argument("--mode", choices=["smoke", "manifest"], default="smoke") # smoke-mode args diff --git a/pyproject.toml b/pyproject.toml index 65b9dfb..054dfcf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mito-train" -version = "0.3.0" +version = "0.3.1" description = "ぴよ将棋OCR 3モデル (board-detector / piece-classifier / hand-classifier) の学習・ONNX出力" requires-python = ">=3.11" dependencies = [ @@ -15,6 +15,7 @@ dependencies = [ "onnxscript>=0.7.1", "opencv-python-headless>=5.0.0.93", "pillow>=12.3.0", + "python-dotenv>=1.1.1", "python-shogi>=1.1.1", "pyyaml>=6.0.3", "timm>=1.0.28", diff --git a/uv.lock b/uv.lock index 348ec64..ac793db 100644 --- a/uv.lock +++ b/uv.lock @@ -1119,7 +1119,7 @@ wheels = [ [[package]] name = "mito-train" -version = "0.3.0" +version = "0.3.1" source = { editable = "." } dependencies = [ { name = "albumentations" }, @@ -1134,6 +1134,7 @@ dependencies = [ { name = "onnxscript" }, { name = "opencv-python-headless" }, { name = "pillow" }, + { name = "python-dotenv" }, { name = "python-shogi" }, { name = "pyyaml" }, { name = "timm" }, @@ -1160,6 +1161,7 @@ requires-dist = [ { name = "onnxscript", specifier = ">=0.7.1" }, { name = "opencv-python-headless", specifier = ">=5.0.0.93" }, { name = "pillow", specifier = ">=12.3.0" }, + { name = "python-dotenv", specifier = ">=1.1.1" }, { name = "python-shogi", specifier = ">=1.1.1" }, { name = "pyyaml", specifier = ">=6.0.3" }, { name = "timm", specifier = ">=1.0.28" }, @@ -2269,6 +2271,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/82/ed/0301aeeac3e5353ef3d94b6ec08bbcabd04a72018415dcb29e588514bba8/python_dotenv-1.2.2.tar.gz", hash = "sha256:2c371a91fbd7ba082c2c1dc1f8bf89ca22564a087c2c287cd9b662adde799cf3", size = 50135, upload-time = "2026-03-01T16:00:26.196Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101, upload-time = "2026-03-01T16:00:25.09Z" }, +] + [[package]] name = "python-shogi" version = "1.1.1" From 7d9526b3f26ac6589b5d72feaec404a91c8d2776 Mon Sep 17 00:00:00 2001 From: tkgstrator Date: Mon, 13 Jul 2026 13:30:35 +0900 Subject: [PATCH 4/6] fix(train): load_dotenv override=True so empty devcontainer forwards lose to .env (#11) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit devcontainer.json forwards ${localEnv:HF_TOKEN} and ${localEnv:WANDB_API_KEY} from the host, which expand to an empty string when those vars are unset on the host. load_dotenv() with the default override=False then treats the empty string as "already set" and refuses to import the value from .env, leaving HF (unauthenticated) and wandb (skipped) silently broken. Switch the three training entrypoints (train_board_ocr, train_piece, train_detector) to load_dotenv(override=True) so .env wins over empty forwards. Real host values still flow through — override just replaces empties too. Bump version to 0.3.2. Co-authored-by: Claude Opus 4.8 (1M context) --- mito_train/training/train_board_ocr.py | 4 +++- mito_train/training/train_detector.py | 4 +++- mito_train/training/train_piece.py | 4 +++- pyproject.toml | 2 +- uv.lock | 2 +- 5 files changed, 11 insertions(+), 5 deletions(-) diff --git a/mito_train/training/train_board_ocr.py b/mito_train/training/train_board_ocr.py index 07b8385..5f9539f 100644 --- a/mito_train/training/train_board_ocr.py +++ b/mito_train/training/train_board_ocr.py @@ -510,7 +510,9 @@ def log_main(msg: str) -> None: def main() -> None: # Pull HF_TOKEN / WANDB_API_KEY / CF_* out of .env into os.environ before # any HF or wandb call resolves credentials. No-op if .env is missing. - load_dotenv() + # override=True so devcontainer.json's ${localEnv:...} forwards that expand + # to an empty string on hosts without those vars don't win over .env. + load_dotenv(override=True) p = argparse.ArgumentParser() p.add_argument("--mode", choices=["smoke", "full"], default="smoke") p.add_argument("--backbone", default="mobilenet_v3_small", diff --git a/mito_train/training/train_detector.py b/mito_train/training/train_detector.py index cd0d953..c539c14 100644 --- a/mito_train/training/train_detector.py +++ b/mito_train/training/train_detector.py @@ -78,7 +78,9 @@ def evaluate(model: nn.Module, loader: DataLoader, device: str) -> dict[str, flo def main() -> None: # Pull HF_TOKEN / WANDB_API_KEY / CF_* out of .env into os.environ before # any wandb or HF call resolves credentials. No-op if .env is missing. - load_dotenv() + # override=True so devcontainer.json's ${localEnv:...} forwards that expand + # to an empty string on hosts without those vars don't win over .env. + load_dotenv(override=True) p = argparse.ArgumentParser() p.add_argument("--train-manifest", type=Path, default=Path("data/detector/train.jsonl")) p.add_argument("--val-manifest", type=Path, default=Path("data/detector/val.jsonl")) diff --git a/mito_train/training/train_piece.py b/mito_train/training/train_piece.py index dd6a5b7..a94013e 100644 --- a/mito_train/training/train_piece.py +++ b/mito_train/training/train_piece.py @@ -223,7 +223,9 @@ def run_manifest(args: argparse.Namespace) -> None: def main() -> None: # Pull WANDB_API_KEY / CF_* / HF_TOKEN out of .env into os.environ before # any wandb or HF call resolves credentials. No-op if .env is missing. - load_dotenv() + # override=True so devcontainer.json's ${localEnv:...} forwards that expand + # to an empty string on hosts without those vars don't win over .env. + load_dotenv(override=True) p = argparse.ArgumentParser() p.add_argument("--mode", choices=["smoke", "manifest"], default="smoke") # smoke-mode args diff --git a/pyproject.toml b/pyproject.toml index 054dfcf..682df05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mito-train" -version = "0.3.1" +version = "0.3.2" description = "ぴよ将棋OCR 3モデル (board-detector / piece-classifier / hand-classifier) の学習・ONNX出力" requires-python = ">=3.11" dependencies = [ diff --git a/uv.lock b/uv.lock index ac793db..1b87464 100644 --- a/uv.lock +++ b/uv.lock @@ -1119,7 +1119,7 @@ wheels = [ [[package]] name = "mito-train" -version = "0.3.1" +version = "0.3.2" source = { editable = "." } dependencies = [ { name = "albumentations" }, From 198dee3daf7762ef8fb7512be98978a1c0f2ada7 Mon Sep 17 00:00:00 2001 From: tkgstrator Date: Mon, 13 Jul 2026 13:34:09 +0900 Subject: [PATCH 5/6] feat(train): add AUTO_FREE_GPUS mode to train_multi_gpu.sh (#14) Automatically restrict CUDA_VISIBLE_DEVICES to GPUs with used memory below FREE_GPU_MEM_MB (default 500 MiB) when AUTO_FREE_GPUS=1. Prevents multiple training jobs from interfering with each other on shared hardware. Co-authored-by: Claude Opus 4.8 (1M context) --- scripts/train_multi_gpu.sh | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/scripts/train_multi_gpu.sh b/scripts/train_multi_gpu.sh index 04360b7..c21c654 100755 --- a/scripts/train_multi_gpu.sh +++ b/scripts/train_multi_gpu.sh @@ -7,11 +7,18 @@ # PER_GPU_WORKERS DataLoader workers per rank (defaults to 4; total across # ranks is NPROC_PER_NODE * PER_GPU_WORKERS) # BACKBONE Model backbone (defaults to mobilenet_v3_small) +# AUTO_FREE_GPUS When 1, restrict training to GPUs whose used memory is +# below FREE_GPU_MEM_MB (default 500). Sets +# CUDA_VISIBLE_DEVICES and NPROC_PER_NODE to the survivors +# so shared boxes don't step on running jobs. +# FREE_GPU_MEM_MB "Free" threshold in MiB (default 500). Only used when +# AUTO_FREE_GPUS=1. # # Examples: # ./scripts/train_multi_gpu.sh # all GPUs, defaults # NPROC_PER_NODE=8 BACKBONE=convnext_tiny ./scripts/train_multi_gpu.sh # NPROC_PER_NODE=4 EPOCHS=100 ./scripts/train_multi_gpu.sh +# AUTO_FREE_GPUS=1 BACKBONE=mobilenet_v3_small ./scripts/train_multi_gpu.sh set -euo pipefail @@ -20,6 +27,34 @@ detect_gpu_count() { uv run python -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo 1 } +# Select GPU indices whose used memory is under FREE_GPU_MEM_MB. Prints a +# comma-separated list (or an empty string if none qualify). +detect_free_gpus() { + local threshold="${1:-500}" + nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits 2>/dev/null \ + | awk -F, -v t="${threshold}" '{ + gsub(/ /, "", $1); gsub(/ /, "", $2); + if ($2+0 < t+0) picks[n++] = $1 + } + END { + for (i=0; i&2 + nvidia-smi --query-gpu=index,memory.used --format=csv >&2 + exit 1 + fi + export CUDA_VISIBLE_DEVICES="${free_list}" + # torchrun sees the remapped devices, so nproc is just the count. + NPROC_PER_NODE=$(awk -F, '{print NF}' <<<"${free_list}") + echo "[train_multi_gpu] AUTO_FREE_GPUS: CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} (threshold=${FREE_GPU_MEM_MB}MiB)" +fi + NPROC_PER_NODE=${NPROC_PER_NODE:-$(detect_gpu_count)} MASTER_PORT=${MASTER_PORT:-29500} PER_GPU_WORKERS=${PER_GPU_WORKERS:-4} From e1e3390bf40f48b341ad0013b77b69321067be73 Mon Sep 17 00:00:00 2001 From: tkgstrator Date: Mon, 13 Jul 2026 16:28:07 +0900 Subject: [PATCH 6/6] feat(train): add cosine LR scheduler with linear warmup (#13) * feat(train): add cosine LR scheduler with linear warmup BoardOCR training previously ran with a constant LR (default 3e-4) for the full epoch budget, leaving cheap image-classification gains on the table especially for 200-epoch runs. Add three CLI flags plumbed into a per-step SequentialLR of LinearLR + CosineAnnealingLR: --scheduler {none,cosine} default: cosine --warmup-epochs INT default: 5 --min-lr-ratio FLOAT default: 0.01 (min_lr = lr * ratio) The scheduler steps once per batch (all ranks share len(train_loader) shard size, so DDP stays in lockstep), its state is persisted in the checkpoint, and current LR is logged to W&B as step/lr each --log-every batch. --scheduler=none reverts to the prior constant-LR behavior for back-compat / A/B comparisons. Note: version bump deferred; PR #11 already bumps to 0.3.2, so whichever PR merges second will need to re-bump. Co-Authored-By: Claude Opus 4.8 (1M context) * chore: bump version to 0.3.3 Companion bump for the LR scheduler feature. develop is at 0.3.2 after PR Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- mito_train/training/train_board_ocr.py | 44 ++++++++++++++++++++++++++ pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 46 insertions(+), 2 deletions(-) diff --git a/mito_train/training/train_board_ocr.py b/mito_train/training/train_board_ocr.py index 5f9539f..f6bd809 100644 --- a/mito_train/training/train_board_ocr.py +++ b/mito_train/training/train_board_ocr.py @@ -26,6 +26,7 @@ import torch.nn as nn import torch.nn.functional as F from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from tqdm.auto import tqdm @@ -282,6 +283,34 @@ def log_main(msg: str) -> None: model.parameters(), lr=args.lr, weight_decay=1e-4, fused=use_cuda, ) + # Per-step scheduler: linear warmup for --warmup-epochs, then cosine anneal + # to lr * --min-lr-ratio over the remaining steps. Every rank steps in + # lockstep (same len(train_loader) shard size), so DDP stays in sync. + steps_per_epoch = len(train_loader) + total_steps = max(1, args.epochs * steps_per_epoch) + warmup_steps = max(0, args.warmup_epochs * steps_per_epoch) + scheduler: torch.optim.lr_scheduler.LRScheduler | None = None + if args.scheduler == "cosine": + cosine_steps = max(1, total_steps - warmup_steps) + eta_min = args.lr * args.min_lr_ratio + if warmup_steps > 0: + warmup = LinearLR( + optimizer, start_factor=1e-3, end_factor=1.0, total_iters=warmup_steps, + ) + cosine = CosineAnnealingLR(optimizer, T_max=cosine_steps, eta_min=eta_min) + scheduler = SequentialLR( + optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps], + ) + else: + scheduler = CosineAnnealingLR(optimizer, T_max=cosine_steps, eta_min=eta_min) + log_main( + f"[board_ocr] scheduler=cosine peak_lr={args.lr:.2e} " + f"warmup={args.warmup_epochs}ep ({warmup_steps} steps) " + f"anneal={cosine_steps} steps -> min_lr={eta_min:.2e}" + ) + else: + log_main(f"[board_ocr] scheduler=none lr={args.lr:.2e} (constant)") + ckpt_dir = args.ckpt_root / args.model_name if is_main(): ckpt_dir.mkdir(parents=True, exist_ok=True) @@ -301,6 +330,8 @@ def log_main(msg: str) -> None: optimizer.load_state_dict(ckpt["optimizer"]) if ckpt.get("scaler") is not None and scaler.is_enabled(): scaler.load_state_dict(ckpt["scaler"]) + if ckpt.get("scheduler") is not None and scheduler is not None: + scheduler.load_state_dict(ckpt["scheduler"]) start_epoch = ckpt["epoch"] + 1 resumed_from = str(ckpt_path) resumed_wandb_id = ckpt.get("wandb_run_id") @@ -329,6 +360,9 @@ def log_main(msg: str) -> None: "effective_batch_size": args.batch_size * world_size, "world_size": world_size, "lr": args.lr, + "scheduler": args.scheduler, + "warmup_epochs": args.warmup_epochs, + "min_lr_ratio": args.min_lr_ratio, "epochs": args.epochs, "start_epoch": start_epoch, "resumed_from": resumed_from, @@ -380,6 +414,8 @@ def log_main(msg: str) -> None: scaler.scale(total).backward() scaler.step(optimizer) scaler.update() + if scheduler is not None: + scheduler.step() running_loss += total.detach() running_board_loss += bl.detach() running_hand_loss += hl.detach() @@ -401,6 +437,7 @@ def log_main(msg: str) -> None: "step/hand_loss": hl.item(), "step/board/cell_acc": m["board/cell_acc"].item(), "step/hand/slot_acc": m["hand/slot_acc"].item(), + "step/lr": optimizer.param_groups[0]["lr"], "step/global_step": global_step, "step/epoch_frac": epoch - 1 + n_batches / len(train_loader), }) @@ -484,6 +521,7 @@ def log_main(msg: str) -> None: "model": _unwrap(model).state_dict(), "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict() if scaler.is_enabled() else None, + "scheduler": scheduler.state_dict() if scheduler is not None else None, "backbone": args.backbone, "image_size": args.image_size, "hand_weight": args.hand_weight, @@ -550,6 +588,12 @@ def main() -> None: 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("--scheduler", choices=["none", "cosine"], default="cosine", + help="Per-step LR schedule. cosine = linear warmup then cosine anneal.") + p.add_argument("--warmup-epochs", type=int, default=5, + help="Epochs of linear warmup (0 -> lr). Ignored when --scheduler=none.") + p.add_argument("--min-lr-ratio", type=float, default=0.01, + help="Cosine anneal floor as a fraction of --lr (min_lr = lr * ratio).") p.add_argument("--hand-mode", choices=["classification", "regression"], default="classification", help="classification: 14x19 logits + CE. regression: 14 scalars + SmoothL1.") diff --git a/pyproject.toml b/pyproject.toml index 682df05..5aadc78 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mito-train" -version = "0.3.2" +version = "0.3.3" description = "ぴよ将棋OCR 3モデル (board-detector / piece-classifier / hand-classifier) の学習・ONNX出力" requires-python = ">=3.11" dependencies = [ diff --git a/uv.lock b/uv.lock index 1b87464..7374b09 100644 --- a/uv.lock +++ b/uv.lock @@ -1119,7 +1119,7 @@ wheels = [ [[package]] name = "mito-train" -version = "0.3.2" +version = "0.3.3" source = { editable = "." } dependencies = [ { name = "albumentations" },