diff --git a/configs/cv_effb0_with_mild.yaml b/configs/cv_effb0_with_mild.yaml new file mode 100644 index 0000000..d59a714 --- /dev/null +++ b/configs/cv_effb0_with_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり込み）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」）．モデル比較用（EfficientNet-B0．モデル名以外は cv_with_mild.yaml と同一） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, mild, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: efficientnet_b0 + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_effb0_without_mild.yaml b/configs/cv_effb0_without_mild.yaml new file mode 100644 index 0000000..853a5bc --- /dev/null +++ b/configs/cv_effb0_without_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり除外．held_out は全程度で評価する）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」）．モデル比較用（EfficientNet-B0．モデル名以外は cv_without_mild.yaml と同一） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: efficientnet_b0 + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_resnet50_with_mild.yaml b/configs/cv_resnet50_with_mild.yaml new file mode 100644 index 0000000..cdb8962 --- /dev/null +++ b/configs/cv_resnet50_with_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり込み）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」）．モデル比較用（ResNet-50．モデル名以外は cv_with_mild.yaml と同一） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, mild, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: resnet50 + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_resnet50_without_mild.yaml b/configs/cv_resnet50_without_mild.yaml new file mode 100644 index 0000000..d696cd9 --- /dev/null +++ b/configs/cv_resnet50_without_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり除外．held_out は全程度で評価する）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」）．モデル比較用（ResNet-50．モデル名以外は cv_without_mild.yaml と同一） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: resnet50 + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_synthetic.yaml b/configs/cv_synthetic.yaml new file mode 100644 index 0000000..ed53684 --- /dev/null +++ b/configs/cv_synthetic.yaml @@ -0,0 +1,27 @@ +# 手元確認用: 合成データ（scripts/make_synthetic_dataset.py の出力）で交差検証を通す設定 +# CPU で数分以内に終わるよう，小さなモデル・事前学習なし・2 エポックにしている． +# data.root の相対パスはカレントディレクトリ（リポジトリのルート）基準． +data: + root: data/synthetic + manifest: manifest.csv # data.root からの相対パス + severities: [strong, mild, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: resnet18 + pretrained: false + +train: + image_size: 64 + batch_size: 8 + max_epochs: 2 + patience: 8 + lr: 1.0e-3 + weight_decay: 0.05 + warmup_epochs: 1 + inner_val_ratio: 0.15 + num_workers: 0 + amp: true diff --git a/configs/cv_vits_with_mild.yaml b/configs/cv_vits_with_mild.yaml new file mode 100644 index 0000000..8e037a3 --- /dev/null +++ b/configs/cv_vits_with_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり込み）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」）．モデル比較用（ViT-Small．モデル名以外は cv_with_mild.yaml と同一） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, mild, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: vit_small_patch16_224 + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_vits_without_mild.yaml b/configs/cv_vits_without_mild.yaml new file mode 100644 index 0000000..f841c28 --- /dev/null +++ b/configs/cv_vits_without_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり除外．held_out は全程度で評価する）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」）．モデル比較用（ViT-Small．モデル名以外は cv_without_mild.yaml と同一） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: vit_small_patch16_224 + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_with_mild.yaml b/configs/cv_with_mild.yaml new file mode 100644 index 0000000..4f568e2 --- /dev/null +++ b/configs/cv_with_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり込み）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, mild, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: convnext_tiny + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git a/configs/cv_without_mild.yaml b/configs/cv_without_mild.yaml new file mode 100644 index 0000000..bc91e04 --- /dev/null +++ b/configs/cv_without_mild.yaml @@ -0,0 +1,27 @@ +# 5-fold 交差検証（ややあり除外．held_out は全程度で評価する）: 計算サーバ上の実データで学習・評価する（SPEC_02「交差検証の実行」） +# 臨床データはサーバ上にのみ置く．手元ではこの設定で実行しない． +# 環境変数 TIAS_DATA_ROOT が設定されていれば data.root より優先される． +data: + root: ~/data/TIASJudgeModel + manifest: manifest.csv # data.root からの相対パス + severities: [strong, normal] + +split: + seed: 42 + n_folds: 5 + +model: + name: convnext_tiny + pretrained: true + +train: + image_size: 224 + batch_size: 32 + max_epochs: 40 + patience: 8 + lr: 1.0e-4 + weight_decay: 0.05 + warmup_epochs: 2 + inner_val_ratio: 0.15 + num_workers: 4 + amp: true diff --git "a/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" "b/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" index bb65820..d46768f 100644 --- "a/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" +++ "b/docs/04_SPEC/SPEC_01_\343\202\242\343\203\274\343\202\255\343\203\206\343\202\257\343\203\201\343\203\243\350\250\255\350\250\210.md" @@ -90,7 +90,7 @@ | `split.n_folds` | 交差検証の fold 数（2 以上） | 5 | - **環境変数 `TIAS_DATA_ROOT` が設定されていれば `data.root` より優先する**（サーバと手元で同じ設定ファイルを使い回せるようにするため） -- 学習・モデル関連の設定は Step 4 で追加する +- 学習・モデル関連の設定は `model`・`train` 節に置く（キーと既定値は [SPEC_02「学習の設定」](SPEC_02_学習・評価仕様.md)）．学習しないコマンド（`tias-dataset`）では省略できる ## クラス定義 (Class Definition) @@ -148,6 +148,7 @@ | コマンド | 役割 | 実行場所 | | --- | --- | --- | | `tias-dataset build-manifest` | `images/` を走査して `manifest.csv` を書き出す（既存があれば `--overwrite` が必要） | サーバ | +| `tias-train cv --config PATH` | 5-fold 交差検証で学習・評価し，`runs/<実験ID>/` に書き出す（[SPEC_02「交差検証の実行」](SPEC_02_学習・評価仕様.md)）．最後に実験 ID を出力する | サーバ（GPU） | | `tias-dataset survey` | データセットを集計して `runs/<実験ID>/` に集計値だけを書き出す（下記） | サーバ（結果の要約のみ手元へ） | | 学習 | 設定 YAML を受け取って学習する | サーバ | | 評価 | 学習済み run の評価・レポート生成 | サーバ | diff --git "a/docs/04_SPEC/SPEC_02_\345\255\246\347\277\222\343\203\273\350\251\225\344\276\241\344\273\225\346\247\230.md" "b/docs/04_SPEC/SPEC_02_\345\255\246\347\277\222\343\203\273\350\251\225\344\276\241\344\273\225\346\247\230.md" index 9f1147b..cc134c2 100644 --- "a/docs/04_SPEC/SPEC_02_\345\255\246\347\277\222\343\203\273\350\251\225\344\276\241\344\273\225\346\247\230.md" +++ "b/docs/04_SPEC/SPEC_02_\345\255\246\347\277\222\343\203\273\350\251\225\344\276\241\344\273\225\346\247\230.md" @@ -30,6 +30,78 @@ まず 1 本を通すことを優先し，比較は timm での差し替えで後から行う． +### 学習の設定 (Training Settings) + +設定 YAML の `model`・`train` 節で与える（[SPEC_01](SPEC_01_アーキテクチャ設計.md)「設定ファイル」）．既定値は次のとおり． + +| 項目 | 既定値 | 備考 | +| --- | --- | --- | +| `model.name` | `convnext_tiny` | timm のモデル名 | +| `model.pretrained` | `true` | ImageNet 事前学習の重み（初回はサーバがダウンロードする）．合成データの試験では `false` | +| `train.image_size` | 224 | アスペクト比を保って長辺を合わせ，余白を黒で埋める（レターボックス） | +| `train.batch_size` | 32 | — | +| `train.max_epochs` | 40 | — | +| `train.patience` | 8 | 検証の macro F1 がこのエポック数改善しなければ打ち切る | +| `train.lr` | 1e-4 | AdamW | +| `train.weight_decay` | 0.05 | AdamW | +| `train.warmup_epochs` | 2 | 線形 warmup の後 cosine で減衰 | +| `train.inner_val_ratio` | 0.15 | 交差検証の各 fold で，学習側から早期終了用の検証データを取る割合（下記） | +| `train.num_workers` | 4 | DataLoader | +| `train.amp` | `true` | CUDA のときだけ混合精度 | + +- データ拡張（学習時のみ）: 左右反転（確率 0.5）・回転 ±10°・拡大縮小 0.9〜1.1 倍・平行移動わずか．はみ出した部分は黒で埋める．**色を変える拡張は行わない**（上記「データ拡張」） +- 正規化: ImageNet の mean / std +- デバイス: CUDA があれば `cuda`（`CUDA_VISIBLE_DEVICES` で割り当てられた GPU），無ければ CPU +- 高速化のため，レターボックス済みの画像（uint8）をメモリに載せてから学習してよい（1085 枚 × 224² × 3 ≒ 160 MB） + +## 交差検証の実行 (Cross-Validation Run) + +**1 つのジョブで 5 つの fold を順に学習・評価**し，1 つの実験ディレクトリ `runs/<実験ID>/` にまとめて書き出す．「ややあり込み」と「ややあり除外」は別のジョブ（別の実験 ID）として回す． + +1. manifest 全体で患者単位・9 グループ層化の 5-fold 分割を作る（`data.severities` に関わらず同じ分割．上記「程度（ややあり）の扱い」） +2. 各 fold について + 1. 学習側（held_out 以外）から `train.inner_val_ratio` の患者を検証用に取る（患者単位・9 グループ層化．seed は分割と同じ設定から決める） + 2. `data.severities` を**学習用と検証用にだけ**適用する（held_out には適用しない．下記） + 3. 検証の macro F1 が最良のエポックの重みを採用する（早期終了） + 4. **held_out（全程度を含む）で 1 回だけ評価する** +3. 全 fold の held_out の予測をまとめて集計する + +### 評価の出し方 (Evaluation Output) + +held_out は `data.severities` に関わらず**全程度の画像**で評価し，次の部分集合ごとに指標を出す．これで 2 つの条件のジョブを同じ土俵で比べられる． + +| 部分集合 | 内容 | 用途 | +| --- | --- | --- | +| `all` | held_out の全画像 | ややあり込みの全体性能 | +| `common` | held_out のうち `strong` と `normal` | **2 条件の比較**（両条件とも学習に含まれる部分） | +| `mild` | held_out のうち `mild` | 境界例の扱い | + +- 各部分集合で: macro F1・accuracy・クラス別 precision / recall / F1・混同行列（`common` と `all`） +- `mild` では加えて**予測の内訳**を出す: 淡紅と予測された割合・元のクラスと予測された割合・その他のクラスと予測された割合（クラス別と全体） +- fold ごとの値と，fold 間の平均 ± 標準偏差を出す．混同行列は全 fold の合計も出す +- 画像ごとの予測（個票）は `predictions.csv` に保存し，**回収しない** + +### 出力 (Run Artifacts) + +[SPEC_01「実験の出力」](SPEC_01_アーキテクチャ設計.md)の名前に従う． + +| ファイル | 内容 | 回収 | +| --- | --- | --- | +| `config.yaml` | 解決済み設定・git ハッシュ・seed | する | +| `metrics.json` | fold ごと・fold 平均の指標（上記の部分集合ごと），各 fold の最良エポック，**所要時間**（下記） | する | +| `history.csv` | fold・エポックごとの学習損失・検証損失・検証 macro F1・学習率・**そのエポックの所要秒数** | する | +| `confusion_matrix.csv` | 全 fold 合計の混同行列（`common` と `all`） | する | +| `log.txt` | ログ | する | +| `figures/*.png` | 学習曲線・混同行列 | する | +| `split.json`・`predictions.csv` | 分割と個票の予測（患者 ID を含む） | しない | +| `checkpoints/fold.pt` | 各 fold の最良の重み | `--with-checkpoints` のときだけ | + +**所要時間の記録**（2026-09-28 ユーザー要望）: `metrics.json` に次を含める．モデル間・条件間で計算コストも比べられるようにするため． + +- 全体の所要秒数（データの読み込みから書き出しまで）と，データの読み込み（レターボックス化）にかかった秒数 +- fold ごとの学習の所要秒数・held_out の評価の所要秒数・実際に回したエポック数・1 エポックあたりの平均秒数 +- 実行環境: デバイス（GPU 名または CPU）・モデル名・パラメータ数 + ## データ分割 (Data Split) **すべての分割は患者単位で行う．** 同一患者の画像が複数の分割に跨ってはならない． diff --git a/pyproject.toml b/pyproject.toml index 7d7d308..edbf234 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ [project.scripts] tias-dataset = "tias_judge_model.cli.dataset:main" tias-remote = "tias_judge_model.cli.remote:main" +tias-train = "tias_judge_model.cli.train:main" [dependency-groups] dev = [ diff --git a/src/tias_judge_model/cli/train.py b/src/tias_judge_model/cli/train.py new file mode 100644 index 0000000..203802d --- /dev/null +++ b/src/tias_judge_model/cli/train.py @@ -0,0 +1,462 @@ +"""`tias-train` コマンド: 患者単位 K-fold 交差検証で学習・評価する（SPEC_01「CLI」）. + +サブコマンド: + +- `cv --config PATH [--out-root runs]`: 5-fold 交差検証で学習・評価し，`runs/<実験ID>/` に書き出す + （SPEC_02「交差検証の実行」）．最後に実験 ID を標準出力へ出す（`tias-remote pull <実験ID>` 用） + +実データを読むためサーバで実行する（手元では合成データの設定 `configs/cv_synthetic.yaml`）． +回収対象（`config.yaml`・`metrics.json`・`history.csv`・`confusion_matrix.csv`・`log.txt`・ +`figures/*.png`）には患者 ID・ファイル名・画像単位の値を出さない．個票（`split.json`・ +`predictions.csv`）はサーバ上にだけ置く． +""" + +from __future__ import annotations + +import argparse +import csv +import json +import logging +import os +import sys +import time +from collections.abc import Mapping, Sequence +from datetime import datetime +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import yaml +from pydantic import ValidationError + +from tias_judge_model.cli.dataset import ( + CONFIG_FILENAME, + FIGURES_DIRNAME, + LOG_FILENAME, + make_run_id, + resolve_git_commit, + resolved_config_document, +) +from tias_judge_model.data.classes import CLASS_SLUGS, label_to_id +from tias_judge_model.data.manifest import ManifestRecord, read_manifest +from tias_judge_model.evaluation.figures import ( + LEARNING_CURVES_FILENAME, + confusion_matrix_filename, + plot_confusion_matrix, + plot_learning_curves, +) +from tias_judge_model.evaluation.metrics import ( + SUBSETS, + aggregate_folds, + mean_std, + subset_metrics, + sum_confusion_matrices, +) +from tias_judge_model.schema import ExperimentConfig, load_config +from tias_judge_model.training.cross_validation import ( + CrossValidationResult, + FoldResult, + fold_seed, + run_cross_validation, +) + +logger = logging.getLogger(__name__) + +METRICS_FILENAME = "metrics.json" +HISTORY_FILENAME = "history.csv" +CONFUSION_MATRIX_FILENAME = "confusion_matrix.csv" +SPLIT_FILENAME = "split.json" +PREDICTIONS_FILENAME = "predictions.csv" +CHECKPOINTS_DIRNAME = "checkpoints" +# 混同行列（全 fold 合計）を書き出す部分集合（SPEC_02「評価の出し方」）． +CONFUSION_SUBSETS: tuple[str, ...] = ("common", "all") + +_LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s" + +_EPILOG = """\ +使用例（サーバ上で実行する．手元では configs/cv_synthetic.yaml で合成データを使う）: + uv run tias-train cv --config configs/cv_with_mild.yaml + uv run tias-remote run --name cv-with-mild -- tias-train cv --config configs/cv_with_mild.yaml + +環境変数: + TIAS_DATA_ROOT data.root を上書き（tias-remote run が各サーバの data_dir を渡す） + TIAS_GIT_COMMIT 実験 ID に入れる git の短縮ハッシュ（tias-remote run が渡す．無ければ nogit） +""" + + +# --------------------------------------------------------------------------- +# cv +# --------------------------------------------------------------------------- + + +def _cmd_cv(args: argparse.Namespace) -> int: + started = time.monotonic() + config = load_config(args.config) + git_commit = resolve_git_commit(os.environ) + run_id = make_run_id(datetime.now(), args.config.stem, git_commit) + run_dir = args.out_root / run_id + if run_dir.exists(): + logger.error("実験ディレクトリが既にある: %s（時間をおいて再実行する）", run_dir) + return 1 + run_dir.mkdir(parents=True) + + file_handler = logging.FileHandler(run_dir / LOG_FILENAME, encoding="utf-8") + file_handler.setFormatter(logging.Formatter(_LOG_FORMAT)) + root = logging.getLogger() + root.addHandler(file_handler) + try: + _run_cv(config, run_dir, run_id=run_id, git_commit=git_commit, started=started) + except (ValueError, OSError) as exc: + # log.txt にも失敗の理由を残すため，ハンドラを外す前にここで記録する． + logger.error("%s", exc) + return 1 + finally: + root.removeHandler(file_handler) + file_handler.close() + sys.stdout.write(f"{run_id}\n") + sys.stdout.flush() + return 0 + + +def _run_cv( + config: ExperimentConfig, run_dir: Path, *, run_id: str, git_commit: str, started: float +) -> None: + logger.info("交差検証の実験 ID: %s", run_id) + document = resolved_config_document( + config, run_id=run_id, git_commit=git_commit, command="train cv" + ) + document["seeds"] = { + "split": config.split.seed, + "folds": [fold_seed(config.split.seed, i) for i in range(config.split.n_folds)], + } + _write_yaml(run_dir / CONFIG_FILENAME, document) + + records = read_manifest(config.data.manifest_path) + result = run_cross_validation(config, records, run_dir / CHECKPOINTS_DIRNAME) + + _write_split(run_dir / SPLIT_FILENAME, config, result) + _write_predictions(run_dir / PREDICTIONS_FILENAME, records, result) + history = history_rows(result.folds) + _write_csv(run_dir / HISTORY_FILENAME, history) + + metrics = build_metrics(config, records, result, run_id=run_id, git_commit=git_commit) + _write_confusion_csv(run_dir / CONFUSION_MATRIX_FILENAME, metrics["confusion_matrix_sum"]) + figures_dir = run_dir / FIGURES_DIRNAME + plot_learning_curves(history, figures_dir / LEARNING_CURVES_FILENAME) + for subset in CONFUSION_SUBSETS: + matrix = metrics["confusion_matrix_sum"].get(subset) + if matrix is not None: + plot_confusion_matrix( + matrix, + f"{subset} (sum of {len(result.folds)} folds held-out)", + figures_dir / confusion_matrix_filename(subset), + ) + + total_seconds = time.monotonic() - started + metrics["timing"]["total_seconds"] = round(total_seconds, 2) + _write_json(run_dir / METRICS_FILENAME, metrics) + + for subset in SUBSETS: + summary = metrics["mean_std"].get(subset) + if summary: + logger.info( + "%s: macro F1 %.4f ± %.4f・accuracy %.4f ± %.4f", + subset, + summary["macro_f1"]["mean"], + summary["macro_f1"]["std"], + summary["accuracy"]["mean"], + summary["accuracy"]["std"], + ) + breakdown = (metrics["pooled"].get("mild") or {}).get("prediction_breakdown") + if breakdown: + overall = breakdown["overall"] + logger.info( + "mild の予測の内訳（全 fold 合計 %d 件）: 淡紅 %.3f・元のクラス %.3f・その他 %.3f", + overall["n"], + overall["light_red"], + overall["original"], + overall["other"], + ) + logger.info( + "出力先: %s（全体の所要 %.1f 秒・うち画像の読み込み %.1f 秒）", + run_dir, + total_seconds, + result.data_loading_seconds, + ) + + +def history_rows(folds: Sequence[FoldResult]) -> list[dict[str, Any]]: + """`history.csv` の行（fold・エポックごとの損失・検証 macro F1・学習率・所要秒数）を作る. + + Args: + folds: fold ごとの結果． + + Returns: + 行の dict のリスト． + """ + rows: list[dict[str, Any]] = [] + for fold in folds: + for h in fold.history: + rows.append( + { + "fold": fold.index, + "epoch": h.epoch, + "train_loss": round(h.train_loss, 6), + "val_loss": round(h.val_loss, 6), + "val_macro_f1": round(h.val_macro_f1, 6), + "lr": h.lr, + "epoch_seconds": round(h.seconds, 3), + } + ) + return rows + + +def build_metrics( + config: ExperimentConfig, + records: Sequence[ManifestRecord], + result: CrossValidationResult, + *, + run_id: str, + git_commit: str, +) -> dict[str, Any]: + """`metrics.json` の内容（fold ごと・fold 平均 ± 標準偏差・全 fold 合算の指標と所要時間）を作る. + + 集計値だけを含み，患者 ID・ファイル名・画像単位の値は含まない． + + Args: + config: 実験設定． + records: manifest の全レコード． + result: 交差検証の結果． + run_id: 実験 ID． + git_commit: git の短縮ハッシュ（または `nogit`）． + + Returns: + JSON に書き出せる dict（`timing.total_seconds` は呼び出し側で埋める）． + """ + labels = np.array([label_to_id(r.label) for r in records], dtype=np.int64) + severities = np.array([r.severity for r in records]) + + fold_entries: list[dict[str, Any]] = [] + all_true: list[np.ndarray] = [] + all_pred: list[np.ndarray] = [] + all_sev: list[np.ndarray] = [] + for fold in result.folds: + idx = fold.held_out_indices + y_true, y_pred, sev = labels[idx], fold.held_out_pred, severities[idx] + all_true.append(y_true) + all_pred.append(y_pred) + all_sev.append(sev) + seconds = [h.seconds for h in fold.history] + fold_entries.append( + { + "fold": fold.index, + "seed": fold.seed, + "best_epoch": fold.best_epoch, + "best_val_macro_f1": fold.best_val_macro_f1, + "counts": { + "patients": { + "train": len(fold.train_patients), + "val": len(fold.val_patients), + "held_out": len(fold.held_out_patients), + }, + "images_by_group": fold.group_counts, + }, + "timing": { + "train_seconds": round(fold.train_seconds, 2), + "eval_seconds": round(fold.eval_seconds, 2), + "epochs_run": len(fold.history), + "mean_epoch_seconds": round(float(np.mean(seconds)), 3) if seconds else None, + }, + "subsets": subset_metrics(y_true, y_pred, sev), + } + ) + + pooled = subset_metrics( + np.concatenate(all_true), np.concatenate(all_pred), np.concatenate(all_sev) + ) + mean_std_by_subset = { + subset: aggregate_folds([f["subsets"][subset] for f in fold_entries]) for subset in SUBSETS + } + confusion_sum = { + subset: ( + sum_confusion_matrices( + [ + f["subsets"][subset]["confusion_matrix"] + for f in fold_entries + if f["subsets"][subset] + ] + ) + if pooled.get(subset) + else None + ) + for subset in CONFUSION_SUBSETS + } + return { + "run_id": run_id, + "git_commit": git_commit, + "classes": list(CLASS_SLUGS), + "severities_trained": list(config.data.severities), + "subsets": {name: list(sevs) for name, sevs in SUBSETS.items()}, + "n_folds": len(result.folds), + "best_epoch": mean_std([f.best_epoch for f in result.folds]), + "mean_std": mean_std_by_subset, + "pooled": pooled, + "confusion_matrix_sum": confusion_sum, + "folds": fold_entries, + "timing": { + "total_seconds": None, + "data_loading_seconds": round(result.data_loading_seconds, 2), + "train_seconds_total": round(sum(f.train_seconds for f in result.folds), 2), + "eval_seconds_total": round(sum(f.eval_seconds for f in result.folds), 2), + }, + "environment": { + "device": result.device, + "model_name": config.model.name, + "pretrained": config.model.pretrained, + "num_parameters": result.num_parameters, + "image_size": config.train.image_size, + "torch_version": torch.__version__, + }, + } + + +def _write_split(path: Path, config: ExperimentConfig, result: CrossValidationResult) -> None: + # 患者 ID を含む（回収しない）． + _write_json( + path, + { + "seed": config.split.seed, + "n_folds": config.split.n_folds, + "inner_val_ratio": config.train.inner_val_ratio, + "severities_trained": list(config.data.severities), + "folds": [ + { + "fold": f.index, + "seed": f.seed, + "train": list(f.train_patients), + "val": list(f.val_patients), + "held_out": list(f.held_out_patients), + } + for f in result.folds + ], + }, + ) + + +def _write_predictions( + path: Path, records: Sequence[ManifestRecord], result: CrossValidationResult +) -> None: + # 画像ごとの予測（個票．患者 ID・ファイル名を含む．回収しない）． + rows: list[dict[str, Any]] = [] + for fold in result.folds: + for i, probs in zip(fold.held_out_indices, fold.held_out_probs, strict=True): + record = records[int(i)] + row: dict[str, Any] = { + "fold": fold.index, + "patient_id": record.patient_id, + "image_path": record.image_path, + "label": record.label, + "severity": record.severity, + "pred": CLASS_SLUGS[int(np.argmax(probs))], + } + row.update( + { + f"prob_{slug}": round(float(p), 6) + for slug, p in zip(CLASS_SLUGS, probs, strict=True) + } + ) + rows.append(row) + _write_csv(path, rows) + + +def _write_confusion_csv(path: Path, matrices: Mapping[str, list[list[int]] | None]) -> None: + rows: list[dict[str, Any]] = [] + for subset, matrix in matrices.items(): + if matrix is None: + continue + for slug, counts in zip(CLASS_SLUGS, matrix, strict=True): + row: dict[str, Any] = {"subset": subset, "true": slug} + row.update({f"pred_{p}": c for p, c in zip(CLASS_SLUGS, counts, strict=True)}) + rows.append(row) + _write_csv(path, rows) + + +def _write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None: + with path.open("w", encoding="utf-8", newline="") as f: + if not rows: + return + writer = csv.DictWriter(f, fieldnames=list(rows[0]), lineterminator="\n") + writer.writeheader() + writer.writerows(rows) + + +def _write_json(path: Path, data: Mapping[str, Any]) -> None: + with path.open("w", encoding="utf-8", newline="\n") as f: + json.dump(data, f, ensure_ascii=False, indent=2) + f.write("\n") + + +def _write_yaml(path: Path, data: Mapping[str, Any]) -> None: + with path.open("w", encoding="utf-8", newline="\n") as f: + yaml.safe_dump(dict(data), f, allow_unicode=True, sort_keys=False) + + +# --------------------------------------------------------------------------- +# エントリポイント +# --------------------------------------------------------------------------- + + +def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="tias-train", + description="患者単位 K-fold 交差検証で学習・評価する", + epilog=_EPILOG, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + sub = parser.add_subparsers(dest="subcommand", required=True) + p_cv = sub.add_parser( + "cv", + help="K-fold 交差検証で学習・評価して runs/<実験ID>/ に書き出す", + description=( + "manifest 全体で患者単位・クラス × 程度の層化の K-fold 分割を作り，各 fold で学習し" + "held_out（全程度）で評価する．最後に実験 ID を標準出力へ出す．" + ), + ) + p_cv.add_argument("--config", type=Path, required=True, help="実験設定 YAML．") + p_cv.add_argument( + "--out-root", type=Path, default=Path("runs"), help="実験ディレクトリの親（既定: runs）．" + ) + return parser.parse_args(argv) + + +def _configure_logging() -> None: + # ログは標準エラーへ（標準出力は実験 ID の受け渡しに使うため汚さない）． + handler = logging.StreamHandler(sys.stderr) + handler.setFormatter(logging.Formatter(_LOG_FORMAT)) + logging.basicConfig(level=logging.INFO, handlers=[handler]) + + +def main(argv: Sequence[str] | None = None) -> int: + """tias-train コマンドのメインエントリポイント. + + Args: + argv: 引数リスト．None のとき sys.argv[1:] を使用． + + Returns: + 終了コード（成功 0・入力の問題 1）． + """ + args = _parse_args(argv) + _configure_logging() + try: + return _cmd_cv(args) + except ValidationError as exc: + logger.error("設定が不正である:\n%s", exc) + return 1 + except (ValueError, OSError) as exc: + logger.error("%s", exc) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/tias_judge_model/data/dataset.py b/src/tias_judge_model/data/dataset.py new file mode 100644 index 0000000..72614c1 --- /dev/null +++ b/src/tias_judge_model/data/dataset.py @@ -0,0 +1,91 @@ +"""学習・評価用の Dataset と入力の正規化（SPEC_02「データ拡張」「前処理」）. + +Dataset はレターボックス済みの uint8 画像を返し，正規化（ImageNet の mean / std）はデバイス上で +`normalize_batch` が行う（CPU のワーカーで float 化するより転送量が小さいため）． + +データ拡張は学習時のみ・幾何変換のみ（左右反転・回転・拡大縮小・平行移動．はみ出しは黒）． +舌の色そのものがラベルであるため，**色を変える拡張は行わない**． +""" + +from __future__ import annotations + +import numpy as np +import torch +from torch.utils.data import Dataset +from torchvision.transforms import v2 + +IMAGENET_MEAN: tuple[float, float, float] = (0.485, 0.456, 0.406) +IMAGENET_STD: tuple[float, float, float] = (0.229, 0.224, 0.225) + +# 幾何拡張の強さ（SPEC_02「学習の設定」の箇条書き）． +_FLIP_PROB = 0.5 +_ROTATION_DEGREES = 10.0 +_SCALE_RANGE = (0.9, 1.1) +_TRANSLATE_RATIO = (0.05, 0.05) + + +def make_train_augmentation() -> v2.Compose: + """学習時の幾何拡張（左右反転・回転 ±10°・拡大縮小 0.9〜1.1・平行移動 ±5%）を作る. + + はみ出した部分は黒 (0) で埋める．色を変える変換は含めない． + + Returns: + `(3, H, W)` の uint8 テンソルを受け取り同じ形状で返す変換． + """ + return v2.Compose( + [ + v2.RandomHorizontalFlip(p=_FLIP_PROB), + v2.RandomAffine( + degrees=_ROTATION_DEGREES, + translate=_TRANSLATE_RATIO, + scale=_SCALE_RANGE, + interpolation=v2.InterpolationMode.BILINEAR, + fill=0, + ), + ] + ) + + +class TongueImageDataset(Dataset): + """レターボックス済み画像とクラス ID の Dataset. + + Args: + images: `(N, H, W, 3)` の uint8 配列（RGB）． + labels: `(N,)` の整数配列（クラス ID．`classes.py` の固定 ID）． + augment: True なら学習時の幾何拡張をかける． + """ + + def __init__(self, images: np.ndarray, labels: np.ndarray, *, augment: bool) -> None: + if len(images) != len(labels): + raise RuntimeError("画像とラベルの件数が一致しない") + self._images = images + self._labels = labels.astype(np.int64) + self._augment = make_train_augmentation() if augment else None + + def __len__(self) -> int: + return len(self._labels) + + def __getitem__(self, index: int) -> tuple[torch.Tensor, int]: + """1 件を返す. + + Returns: + (`(3, H, W)` の uint8 テンソル, クラス ID)． + """ + image = torch.from_numpy(np.ascontiguousarray(self._images[index])).permute(2, 0, 1) + if self._augment is not None: + image = self._augment(image) + return image, int(self._labels[index]) + + +def normalize_batch(images: torch.Tensor) -> torch.Tensor: + """uint8 画像のバッチを float にして ImageNet の mean / std で正規化する. + + Args: + images: `(B, 3, H, W)` の uint8 テンソル（任意のデバイス）． + + Returns: + `(B, 3, H, W)` の float32 テンソル（同じデバイス）． + """ + mean = torch.tensor(IMAGENET_MEAN, device=images.device).view(1, 3, 1, 1) + std = torch.tensor(IMAGENET_STD, device=images.device).view(1, 3, 1, 1) + return (images.float() / 255.0 - mean) / std diff --git a/src/tias_judge_model/data/images.py b/src/tias_judge_model/data/images.py new file mode 100644 index 0000000..0e88835 --- /dev/null +++ b/src/tias_judge_model/data/images.py @@ -0,0 +1,78 @@ +"""画像の読み込みとレターボックス化（SPEC_02「前処理」）. + +アスペクト比を保って長辺を `image_size` に合わせ，余白を黒 (0, 0, 0) で埋める．マスク外の +領域も黒であるため，パディングも黒で統一する．学習を速くするため，全画像をレターボックス化した +uint8 配列としてメモリに載せる（1085 枚 × 224² × 3 ≒ 160 MB）． +""" + +from __future__ import annotations + +import logging +from collections.abc import Sequence +from pathlib import Path + +import numpy as np +from PIL import Image + +from tias_judge_model.data.manifest import ManifestRecord, resolve_image_path + +logger = logging.getLogger(__name__) + + +def letterbox(image: Image.Image, size: int) -> np.ndarray: + """画像をアスペクト比を保って `size × size` に収め，余白を黒で埋める. + + 長辺が `size` になるよう縮小（または拡大）し，中央に配置する． + + Args: + image: RGB の PIL 画像． + size: 出力の一辺（ピクセル）． + + Returns: + `(size, size, 3)` の uint8 配列（RGB）． + """ + width, height = image.size + scale = size / max(width, height) + new_w = max(1, round(width * scale)) + new_h = max(1, round(height * scale)) + resized = image.resize((new_w, new_h), Image.Resampling.BILINEAR) + canvas = Image.new("RGB", (size, size), (0, 0, 0)) + canvas.paste(resized, ((size - new_w) // 2, (size - new_h) // 2)) + return np.asarray(canvas, dtype=np.uint8) + + +def load_letterboxed_images( + records: Sequence[ManifestRecord], images_dir: Path, size: int +) -> np.ndarray: + """manifest の全画像を読み，レターボックス化して 1 つの配列にまとめる. + + 読めない画像は学習を歪めるため止める（fail fast）．例外メッセージには画像のパスを出さず， + manifest のデータ行の番号（1 始まり）で示す． + + Args: + records: manifest レコード． + images_dir: 画像ディレクトリ（`/images`）． + size: 出力の一辺（ピクセル）． + + Returns: + `(N, size, size, 3)` の uint8 配列（RGB）．並びは `records` と同じ． + + Raises: + ValueError: 画像が無い，または読めない場合． + """ + out = np.zeros((len(records), size, size, 3), dtype=np.uint8) + for i, record in enumerate(records): + path = resolve_image_path(images_dir, record) + if not path.is_file(): + raise ValueError(f"manifest のデータ {i + 1} 件目: 画像ファイルが無い") + try: + with Image.open(path) as img: + img.load() + rgb = img.convert("RGB") + except (OSError, ValueError, Image.DecompressionBombError) as exc: + raise ValueError( + f"manifest のデータ {i + 1} 件目: 画像を読めない（{type(exc).__name__}）" + ) from None + out[i] = letterbox(rgb, size) + logger.info("画像 %d 件をレターボックス化した（%d×%d）", len(records), size, size) + return out diff --git a/src/tias_judge_model/data/split.py b/src/tias_judge_model/data/split.py index 48f656c..9560fed 100644 --- a/src/tias_judge_model/data/split.py +++ b/src/tias_judge_model/data/split.py @@ -261,6 +261,53 @@ return folds +def split_inner_val( + records: Sequence[ManifestRecord], val_ratio: float, seed: int +) -> tuple[tuple[str, ...], tuple[str, ...]]: + """交差検証の fold の学習側を，患者単位・クラス × 程度の層化で学習用と検証用に分ける. + + `split_train_val_test` と同じ方式（層化グループごとに患者 ID を昇順に並べてから seed で + 並べ替え，先頭から `round(患者数 × 比率)` 人を検証用に取る）で，test を取らない 2 分割を作る． + `data.severities` による除外はここでは行わない（分割後に適用する）． + + Args: + records: fold の学習側の manifest レコード（除外前）． + val_ratio: 検証用に回す患者の比率（0 より大きく 1 未満）． + seed: 乱数 seed． + + Returns: + (学習用の患者 ID, 検証用の患者 ID)．それぞれ昇順． + + Raises: + ValueError: 比率が不正，レコードが空，同一患者に異なる label がある， + またはある層化グループで学習用に患者が残らない場合． + """ + if not 0.0 < val_ratio < 1.0: + raise ValueError(f"検証用の比率が不正である（{val_ratio}）．0 より大きく 1 未満であること") + by_group = _patients_by_group(records) + rng = np.random.default_rng(seed) + train: list[str] = [] + val: list[str] = [] + for group in ALL_DIR_NAMES: + patients = by_group.get(group) + if not patients: + continue + shuffled = [patients[i] for i in rng.permutation(len(patients))] + n_val = _round_half_up(len(patients) * val_ratio) + if n_val >= len(patients): + raise ValueError( + f"グループ {group} の患者 {len(patients)} 人では学習用に患者が残らない" + f"（検証用 {n_val} 人）" + ) + if n_val == 0: + logger.warning( + "グループ %s は患者 %d 人のため検証用 0 人になった", group, len(patients) + ) + val.extend(shuffled[:n_val]) + train.extend(shuffled[n_val:]) + return tuple(sorted(train)), tuple(sorted(val)) + + def select_records( records: Iterable[ManifestRecord], patient_ids: Collection[str] ) -> list[ManifestRecord]: diff --git a/src/tias_judge_model/evaluation/__init__.py b/src/tias_judge_model/evaluation/__init__.py new file mode 100644 index 0000000..fd78ef3 --- /dev/null +++ b/src/tias_judge_model/evaluation/__init__.py @@ -0,0 +1 @@ +"""評価層: 指標・予測の内訳・fold 間の集計・図.""" diff --git a/src/tias_judge_model/evaluation/figures.py b/src/tias_judge_model/evaluation/figures.py new file mode 100644 index 0000000..4ca78f9 --- /dev/null +++ b/src/tias_judge_model/evaluation/figures.py @@ -0,0 +1,101 @@ +"""交差検証の図（学習曲線・混同行列）. + +`figures/` は回収対象であるため，集計値の図だけを描き，画像そのものは描かない（SPEC_01 +「実験の出力」）．`data/survey_figures.py` と同じく pyplot を使わず `Figure` ＋ +`FigureCanvasAgg` で描く．ラベルは日本語フォントの有無に左右されないよう英字で書く． +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from pathlib import Path + +import numpy as np +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure + +from tias_judge_model.data.classes import CLASS_SLUGS + +LEARNING_CURVES_FILENAME = "learning_curves.png" +_DPI = 120 + + +def confusion_matrix_filename(subset: str) -> str: + """部分集合の混同行列の図のファイル名を返す（例: `confusion_matrix_all.png`）.""" + return f"confusion_matrix_{subset}.png" + + +def plot_learning_curves(history: Sequence[Mapping[str, float]], path: Path) -> None: + """fold ごとの学習曲線（損失・検証 macro F1）を描いて PNG に保存する. + + Args: + history: エポックごとの記録（`fold`・`epoch`・`train_loss`・`val_loss`・`val_macro_f1`）． + path: 出力先の PNG パス．親ディレクトリは作成される． + """ + folds = sorted({int(h["fold"]) for h in history}) + fig = Figure(figsize=(12, 5), dpi=_DPI) + FigureCanvasAgg(fig) + ax_loss, ax_f1 = fig.subplots(1, 2) + for fold in folds: + rows = [h for h in history if int(h["fold"]) == fold] + epochs = [int(h["epoch"]) for h in rows] + line = ax_loss.plot(epochs, [h["train_loss"] for h in rows], label=f"fold{fold} train")[0] + ax_loss.plot( + epochs, + [h["val_loss"] for h in rows], + linestyle="--", + color=line.get_color(), + label=f"fold{fold} val", + ) + ax_f1.plot(epochs, [h["val_macro_f1"] for h in rows], label=f"fold{fold}") + ax_loss.set_xlabel("epoch") + ax_loss.set_ylabel("loss") + ax_loss.set_title("loss (solid: train, dashed: val)") + ax_loss.legend(fontsize=7) + ax_f1.set_xlabel("epoch") + ax_f1.set_ylabel("val macro F1") + ax_f1.set_ylim(0.0, 1.0) + ax_f1.set_title("validation macro F1") + ax_f1.legend(fontsize=7) + fig.tight_layout() + path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(path) + + +def plot_confusion_matrix(matrix: Sequence[Sequence[int]], title: str, path: Path) -> None: + """混同行列（件数と行方向の割合）を描いて PNG に保存する. + + Args: + matrix: 5×5 の混同行列（行が正解・列が予測．クラス ID 順）． + title: 図のタイトル． + path: 出力先の PNG パス．親ディレクトリは作成される． + """ + cm = np.asarray(matrix, dtype=np.float64) + row_sums = cm.sum(axis=1, keepdims=True) + ratio = np.divide(cm, row_sums, out=np.zeros_like(cm), where=row_sums > 0) + fig = Figure(figsize=(7, 6), dpi=_DPI) + FigureCanvasAgg(fig) + ax = fig.subplots() + image = ax.imshow(ratio, vmin=0.0, vmax=1.0, cmap="Blues") + fig.colorbar(image, ax=ax, label="row ratio") + n = len(CLASS_SLUGS) + ax.set_xticks(range(n), CLASS_SLUGS, rotation=30) + ax.set_yticks(range(n), CLASS_SLUGS) + ax.set_xlabel("predicted") + ax.set_ylabel("true") + ax.set_title(title) + for i in range(n): + for j in range(n): + color = "white" if ratio[i, j] > 0.5 else "black" + ax.text( + j, + i, + f"{int(cm[i, j])}\n{ratio[i, j]:.2f}", + ha="center", + va="center", + color=color, + fontsize=8, + ) + fig.tight_layout() + path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(path) diff --git a/src/tias_judge_model/evaluation/metrics.py b/src/tias_judge_model/evaluation/metrics.py new file mode 100644 index 0000000..246ae42 --- /dev/null +++ b/src/tias_judge_model/evaluation/metrics.py @@ -0,0 +1,223 @@ +"""評価指標（SPEC_02「評価指標」「交差検証の実行」）. + +すべて入力 → 出力の純関数．クラスの軸は常に `classes.py` の固定 ID 順（混同行列の軸をぶらさない）． +出力は集計値だけであり，画像単位の値・患者 ID は含まない． +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np +from sklearn.metrics import confusion_matrix, f1_score, precision_recall_fscore_support + +from tias_judge_model.data.classes import CLASS_SLUGS, NUM_CLASSES, label_to_id + +# 淡紅（正常）のクラス ID．mild の予測の内訳で「淡紅と予測された割合」に使う． +LIGHT_RED_ID: int = label_to_id("light_red") + +# 評価の部分集合と含める程度（SPEC_02「評価の出し方」）． +SUBSETS: dict[str, tuple[str, ...]] = { + "all": ("strong", "mild", "normal"), + "common": ("strong", "normal"), + "mild": ("mild",), +} + +# fold 間の平均・標準偏差を取らないキー（件数・行列）． +_NON_AVERAGED_KEYS = frozenset({"n", "support", "confusion_matrix"}) + + +def macro_f1(y_true: np.ndarray, y_pred: np.ndarray) -> float: + """macro F1 を返す（`y_true` に現れるクラスで平均する）. + + `y_true` に現れないクラス（例: mild の部分集合の淡紅）を平均に含めると，そのクラスへの予測が + それだけで F1 = 0 として平均を下げるため，正解側に現れるクラスだけで平均する． + + Args: + y_true: `(N,)` の正解クラス ID． + y_pred: `(N,)` の予測クラス ID． + + Returns: + macro F1（N = 0 なら 0.0）． + """ + if len(y_true) == 0: + return 0.0 + labels = np.unique(y_true) + return float(f1_score(y_true, y_pred, labels=labels, average="macro", zero_division=0)) + + +def classification_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, Any]: + """分類指標（macro F1・accuracy・クラス別 precision / recall / F1・混同行列）を返す. + + Args: + y_true: `(N,)` の正解クラス ID（0〜4）． + y_pred: `(N,)` の予測クラス ID（0〜4）． + + Returns: + `{"n", "macro_f1", "accuracy", "per_class": {slug: {precision, recall, f1, support}}, + "confusion_matrix": [[...]]}`．混同行列は行が正解・列が予測でクラス ID 順（5×5）． + クラス別は `y_true` に現れるクラスだけを含む． + """ + y_true = np.asarray(y_true, dtype=np.int64) + y_pred = np.asarray(y_pred, dtype=np.int64) + all_labels = list(range(NUM_CLASSES)) + cm = confusion_matrix(y_true, y_pred, labels=all_labels) + precision, recall, f1, support = precision_recall_fscore_support( + y_true, y_pred, labels=all_labels, zero_division=0 + ) + present = set(np.unique(y_true).tolist()) + per_class = { + CLASS_SLUGS[c]: { + "precision": float(precision[c]), + "recall": float(recall[c]), + "f1": float(f1[c]), + "support": int(support[c]), + } + for c in all_labels + if c in present + } + return { + "n": int(len(y_true)), + "macro_f1": macro_f1(y_true, y_pred), + "accuracy": float(np.mean(y_true == y_pred)) if len(y_true) else 0.0, + "per_class": per_class, + "confusion_matrix": cm.astype(int).tolist(), + } + + +def mild_prediction_breakdown(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, Any]: + """mild の画像の予測の内訳（淡紅・元のクラス・その他と予測された割合）を返す. + + 淡紅と予測することを誤りと断定せず，割合として報告する（SPEC_02「程度（ややあり）の扱い」）． + 淡紅には mild が無いため，元のクラスと淡紅は重ならない． + + Args: + y_true: `(N,)` の mild の画像の正解クラス ID． + y_pred: `(N,)` の予測クラス ID． + + Returns: + `{"overall": {...}, "per_class": {slug: {...}}}`．各値は + `{"n", "light_red", "original", "other"}`（割合．N = 0 なら None）． + クラス別は `y_true` に現れるクラスだけを含む． + """ + y_true = np.asarray(y_true, dtype=np.int64) + y_pred = np.asarray(y_pred, dtype=np.int64) + per_class = { + CLASS_SLUGS[c]: _breakdown(y_true[y_true == c], y_pred[y_true == c]) + for c in range(NUM_CLASSES) + if np.any(y_true == c) + } + return {"overall": _breakdown(y_true, y_pred), "per_class": per_class} + + +def _breakdown(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, Any]: + n = len(y_true) + if n == 0: + return {"n": 0, "light_red": None, "original": None, "other": None} + to_light_red = y_pred == LIGHT_RED_ID + to_original = (y_pred == y_true) & ~to_light_red + other = ~to_light_red & ~to_original + return { + "n": int(n), + "light_red": float(np.mean(to_light_red)), + "original": float(np.mean(to_original)), + "other": float(np.mean(other)), + } + + +def subset_metrics( + y_true: np.ndarray, y_pred: np.ndarray, severities: Sequence[str] +) -> dict[str, Any]: + """部分集合 `all`・`common`・`mild` ごとの指標を返す（SPEC_02「評価の出し方」）. + + Args: + y_true: `(N,)` の正解クラス ID． + y_pred: `(N,)` の予測クラス ID． + severities: 長さ N の各画像の程度（`strong`・`mild`・`normal`）． + + Returns: + 部分集合名 → `classification_metrics` の結果（該当 0 件なら None）． + `mild` には予測の内訳 `prediction_breakdown` を加える． + """ + y_true = np.asarray(y_true, dtype=np.int64) + y_pred = np.asarray(y_pred, dtype=np.int64) + sev = np.asarray(severities) + result: dict[str, Any] = {} + for name, included in SUBSETS.items(): + mask = np.isin(sev, included) + if not np.any(mask): + result[name] = None + continue + metrics = classification_metrics(y_true[mask], y_pred[mask]) + if name == "mild": + metrics["prediction_breakdown"] = mild_prediction_breakdown(y_true[mask], y_pred[mask]) + result[name] = metrics + return result + + +def mean_std(values: Sequence[float]) -> dict[str, float | int | None]: + """値の平均と標準偏差（不偏．1 件なら 0）を返す. + + Args: + values: fold ごとの値． + + Returns: + `{"mean", "std", "n"}`（0 件なら mean・std は None）． + """ + arr = np.asarray(values, dtype=np.float64) + if arr.size == 0: + return {"mean": None, "std": None, "n": 0} + std = float(np.std(arr, ddof=1)) if arr.size > 1 else 0.0 + return {"mean": float(np.mean(arr)), "std": std, "n": int(arr.size)} + + +def aggregate_folds(fold_values: Sequence[Mapping[str, Any] | None]) -> dict[str, Any] | None: + """fold ごとの入れ子の指標（同じ形の dict）から，数値の葉ごとに平均 ± 標準偏差を作る. + + 件数（`n`・`support`）と混同行列は平均しない．ある fold に無いキー（例: その fold の + mild に現れないクラス）は，そのキーを持つ fold だけで集計する（`n` に fold 数が入る）． + + Args: + fold_values: fold ごとの指標（None の fold は除く）． + + Returns: + 同じ形の dict で，数値の葉を `{"mean", "std", "n"}` に置き換えたもの． + 有効な fold が無ければ None． + """ + present = [v for v in fold_values if v is not None] + if not present: + return None + keys: list[str] = [] + for v in present: + keys.extend(k for k in v if k not in keys) + result: dict[str, Any] = {} + for key in keys: + if key in _NON_AVERAGED_KEYS: + continue + items = [v.get(key) for v in present] + items = [x for x in items if x is not None] + if not items: + continue + if all(isinstance(x, Mapping) for x in items): + nested = aggregate_folds(items) + if nested: + result[key] = nested + elif all(isinstance(x, int | float) and not isinstance(x, bool) for x in items): + result[key] = mean_std([float(x) for x in items]) + return result + + +def sum_confusion_matrices(matrices: Sequence[Sequence[Sequence[int]]]) -> list[list[int]]: + """混同行列（同じ形）の合計を返す. + + Args: + matrices: 混同行列のリスト（各 5×5）． + + Returns: + 合計の混同行列（5×5．空なら 0 行列）． + """ + total = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.int64) + for m in matrices: + total += np.asarray(m, dtype=np.int64) + return total.astype(int).tolist() diff --git a/src/tias_judge_model/models/__init__.py b/src/tias_judge_model/models/__init__.py new file mode 100644 index 0000000..664250e --- /dev/null +++ b/src/tias_judge_model/models/__init__.py @@ -0,0 +1 @@ +"""モデル層: timm のモデル生成（アーキテクチャの差し替え点）.""" diff --git a/src/tias_judge_model/models/factory.py b/src/tias_judge_model/models/factory.py new file mode 100644 index 0000000..66904b9 --- /dev/null +++ b/src/tias_judge_model/models/factory.py @@ -0,0 +1,42 @@ +"""timm によるモデル生成（SPEC_02「モデル」）. + +`model.name` に timm のモデル名を渡すだけでアーキテクチャを差し替えられるようにする． +""" + +from __future__ import annotations + +import timm +from torch import nn + +from tias_judge_model.data.classes import NUM_CLASSES + + +def create_model(name: str, *, pretrained: bool, num_classes: int = NUM_CLASSES) -> nn.Module: + """timm のモデルを出力 `num_classes` クラスで作る. + + Args: + name: timm のモデル名（例: `convnext_tiny`・`efficientnet_b0`・`vit_small_patch16_224`）． + pretrained: ImageNet 事前学習の重みを使うか（初回はダウンロードする）． + num_classes: 出力クラス数． + + Returns: + 入力 `(B, 3, H, W)` の float テンソルから `(B, num_classes)` のロジットを返すモデル． + + Raises: + ValueError: timm に無いモデル名の場合． + """ + if not timm.is_model(name): + raise ValueError(f"timm に無いモデル名である: {name!r}") + return timm.create_model(name, pretrained=pretrained, num_classes=num_classes) + + +def count_parameters(model: nn.Module) -> int: + """モデルのパラメータ数（学習対象外も含む）を返す. + + Args: + model: モデル． + + Returns: + パラメータの要素数の合計． + """ + return sum(p.numel() for p in model.parameters()) diff --git a/src/tias_judge_model/schema.py b/src/tias_judge_model/schema.py index f01df80..313c9ea 100644 --- a/src/tias_judge_model/schema.py +++ b/src/tias_judge_model/schema.py @@ -1,6 +1,7 @@ """pydantic による設定スキーマと設定 YAML の読み込み. -学習・モデル関連の設定は後続のステップで `ExperimentConfig` にフィールドを追加して拡張する． +学習・モデル関連の設定（`model`・`train`）は省略でき，省略時は SPEC_02「学習の設定」の +既定値で補う（学習しないコマンド `tias-dataset` の設定ファイルをそのまま読めるようにするため）． """ from __future__ import annotations @@ -102,16 +103,60 @@ return self +class ModelConfig(_StrictModel): + """モデルの設定（SPEC_02「学習の設定」）. + + Attributes: + name: timm のモデル名（例: `convnext_tiny`・`resnet50`・`vit_small_patch16_224`）． + pretrained: ImageNet 事前学習の重みを使うか． + """ + + name: str = Field(default="convnext_tiny", min_length=1) + pretrained: bool = True + + +class TrainConfig(_StrictModel): + """学習の設定（SPEC_02「学習の設定」）. + + Attributes: + image_size: 入力画像の一辺（長辺を合わせるレターボックス）． + batch_size: バッチサイズ． + max_epochs: 最大エポック数． + patience: 検証 macro F1 がこのエポック数改善しなければ打ち切る． + lr: AdamW の学習率． + weight_decay: AdamW の weight decay． + warmup_epochs: 線形 warmup のエポック数（その後 cosine で減衰）． + inner_val_ratio: 交差検証の各 fold で，学習側から早期終了用の検証に回す患者の比率． + num_workers: DataLoader のワーカー数． + amp: CUDA のときに混合精度を使うか（CPU では無視する）． + """ + + image_size: int = Field(default=224, ge=16) + batch_size: int = Field(default=32, ge=1) + max_epochs: int = Field(default=40, ge=1) + patience: int = Field(default=8, ge=1) + lr: float = Field(default=1e-4, gt=0.0) + weight_decay: float = Field(default=0.05, ge=0.0) + warmup_epochs: int = Field(default=2, ge=0) + inner_val_ratio: float = Field(default=0.15, gt=0.0, lt=1.0) + num_workers: int = Field(default=4, ge=0) + amp: bool = True + + class ExperimentConfig(_StrictModel): """実験設定のルート. Attributes: data: データの場所． split: 分割の設定． + model: モデルの設定（省略時は既定値）． + train: 学習の設定（省略時は既定値）． """ data: DataConfig split: SplitConfig = Field(default_factory=SplitConfig) + model: ModelConfig = Field(default_factory=ModelConfig) + train: TrainConfig = Field(default_factory=TrainConfig) def load_config(path: Path, environ: Mapping[str, str] | None = None) -> ExperimentConfig: diff --git a/src/tias_judge_model/training/__init__.py b/src/tias_judge_model/training/__init__.py new file mode 100644 index 0000000..d96c5d4 --- /dev/null +++ b/src/tias_judge_model/training/__init__.py @@ -0,0 +1 @@ +"""学習層: 学習ループ・交差検証の実行.""" diff --git a/src/tias_judge_model/training/cross_validation.py b/src/tias_judge_model/training/cross_validation.py new file mode 100644 index 0000000..aa80e85 --- /dev/null +++ b/src/tias_judge_model/training/cross_validation.py @@ -0,0 +1,296 @@ +"""患者単位 K-fold 交差検証の実行（SPEC_02「交差検証の実行」）. + +1. manifest 全体で患者単位・9 グループ層化の K-fold 分割を作る（`data.severities` に関わらず同じ） +2. 各 fold で + 1. 学習側から `train.inner_val_ratio` の患者を検証用に取る（患者単位・9 グループ層化） + 2. `data.severities` を学習用と検証用にだけ適用する + 3. 検証 macro F1 が最良のエポックの重みを採用する（早期終了） + 4. held_out（**全程度**）で 1 回だけ評価する +3. 全 fold の held_out の予測を呼び出し側でまとめて集計する + +ログには件数・グループ別件数・損失・指標・時間だけを出し，患者 ID・ファイル名は出さない． +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import DataLoader + +from tias_judge_model.data.classes import label_to_id +from tias_judge_model.data.dataset import TongueImageDataset +from tias_judge_model.data.images import load_letterboxed_images +from tias_judge_model.data.manifest import ManifestRecord +from tias_judge_model.data.split import ( + count_by_group, + filter_by_severity, + make_kfold_splits, + select_records, + split_inner_val, +) +from tias_judge_model.models.factory import count_parameters, create_model +from tias_judge_model.schema import ExperimentConfig, TrainConfig +from tias_judge_model.training.loop import ( + EpochRecord, + device_name, + predict_logits, + seed_everything, + select_device, + train_model, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class FoldResult: + """1 fold の学習・評価の結果. + + `train_patients`・`val_patients`・`held_out_patients` と `held_out_indices` は個票に当たるため， + サーバ上の `split.json`・`predictions.csv` にだけ書き出す． + + Attributes: + index: fold 番号（0 始まり）． + seed: この fold の学習に使った seed． + train_patients: 学習用の患者 ID（除外前．昇順）． + val_patients: 検証用の患者 ID（除外前．昇順）． + held_out_patients: held_out の患者 ID（昇順）． + group_counts: 分割名（`train`・`val`・`held_out`）→ 層化グループ → 画像枚数 + （`train`・`val` は `data.severities` 適用後，`held_out` は全程度）． + history: エポックごとの記録． + best_epoch: 最良エポック（1 始まり）． + best_val_macro_f1: 最良エポックの検証 macro F1． + train_seconds: 学習の所要秒数． + eval_seconds: held_out の評価の所要秒数． + held_out_indices: `(M,)` の held_out 画像の manifest 上の位置（0 始まり）． + held_out_probs: `(M, C)` の held_out の予測確率（softmax）． + """ + + index: int + seed: int + train_patients: tuple[str, ...] + val_patients: tuple[str, ...] + held_out_patients: tuple[str, ...] + group_counts: dict[str, dict[str, int]] + history: list[EpochRecord] + best_epoch: int + best_val_macro_f1: float + train_seconds: float + eval_seconds: float + held_out_indices: np.ndarray + held_out_probs: np.ndarray + + @property + def held_out_pred(self) -> np.ndarray: + """`(M,)` の held_out の予測クラス ID.""" + return self.held_out_probs.argmax(axis=1) + + +@dataclass +class CrossValidationResult: + """交差検証全体の結果. + + Attributes: + folds: fold ごとの結果． + data_loading_seconds: 画像の読み込み（レターボックス化）の所要秒数． + device: デバイスの表示名（GPU 名または `cpu`）． + num_parameters: モデルのパラメータ数． + """ + + folds: list[FoldResult] = field(default_factory=list) + data_loading_seconds: float = 0.0 + device: str = "cpu" + num_parameters: int = 0 + + +def fold_seed(base_seed: int, fold_index: int) -> int: + """fold ごとの seed（内側の検証分割と学習の乱数）を分割の seed から決める. + + Args: + base_seed: 設定の `split.seed`． + fold_index: fold 番号（0 始まり）． + + Returns: + `base_seed + fold_index`． + """ + return base_seed + fold_index + + +def run_cross_validation( + config: ExperimentConfig, records: Sequence[ManifestRecord], checkpoint_dir: Path +) -> CrossValidationResult: + """K-fold 交差検証で学習・評価する（各 fold の最良の重みを `checkpoint_dir` に保存する）. + + Args: + config: 実験設定． + records: manifest の全レコード（検証済み．除外前）． + checkpoint_dir: 各 fold の最良の重み `fold.pt` の保存先． + + Returns: + 交差検証の結果． + + Raises: + ValueError: 分割できない，画像を読めない，学習用・検証用が空になる等の入力の問題． + """ + records = list(records) + device = select_device() + result = CrossValidationResult(device=device_name(device)) + logger.info( + "交差検証を開始する: 画像 %d 件・%d fold・モデル %s（pretrained=%s）・程度 %s・デバイス %s", + len(records), + config.split.n_folds, + config.model.name, + config.model.pretrained, + list(config.data.severities), + result.device, + ) + + folds = make_kfold_splits(records, config.split.n_folds, config.split.seed) + + started = time.monotonic() + images = load_letterboxed_images(records, config.data.images_dir, config.train.image_size) + result.data_loading_seconds = time.monotonic() - started + logger.info("画像の読み込みに %.1f 秒かかった", result.data_loading_seconds) + labels = np.array([label_to_id(r.label) for r in records], dtype=np.int64) + position = {r.image_path: i for i, r in enumerate(records)} + + checkpoint_dir.mkdir(parents=True, exist_ok=True) + n_folds = len(folds) + for fold in folds: + prefix = f"fold{fold.index}（{fold.index + 1}/{n_folds}）" + seed = fold_seed(config.split.seed, fold.index) + train_side = select_records(records, fold.train) + train_patients, val_patients = split_inner_val( + train_side, config.train.inner_val_ratio, seed + ) + train_records = filter_by_severity( + select_records(train_side, train_patients), config.data.severities + ) + val_records = filter_by_severity( + select_records(train_side, val_patients), config.data.severities + ) + # held_out には data.severities を適用しない（全程度で評価する）． + held_records = select_records(records, fold.held_out) + if not train_records or not val_records or not held_records: + raise ValueError( + f"{prefix}: 学習用 {len(train_records)} 件・検証用 {len(val_records)} 件・" + f"held_out {len(held_records)} 件のいずれかが 0 件である" + ) + group_counts = { + "train": count_by_group(train_records), + "val": count_by_group(val_records), + "held_out": count_by_group(held_records), + } + logger.info( + "%s: 画像 学習 %d・検証 %d・held_out %d（患者 学習 %d・検証 %d・held_out %d）", + prefix, + len(train_records), + len(val_records), + len(held_records), + len(train_patients), + len(val_patients), + len(fold.held_out), + ) + for name, counts in group_counts.items(): + logger.info("%s: %s のグループ別枚数 %s", prefix, name, counts) + + train_idx = np.array([position[r.image_path] for r in train_records], dtype=np.int64) + val_idx = np.array([position[r.image_path] for r in val_records], dtype=np.int64) + held_idx = np.array([position[r.image_path] for r in held_records], dtype=np.int64) + + seed_everything(seed) + model = create_model(config.model.name, pretrained=config.model.pretrained).to(device) + result.num_parameters = count_parameters(model) + train_loader = _make_loader( + images[train_idx], labels[train_idx], config.train, device, train=True, seed=seed + ) + val_loader = _make_loader( + images[val_idx], labels[val_idx], config.train, device, train=False, seed=seed + ) + held_loader = _make_loader( + images[held_idx], labels[held_idx], config.train, device, train=False, seed=seed + ) + + trained = train_model( + model, train_loader, val_loader, config.train, device, log_prefix=prefix + ) + torch.save( + { + "fold": fold.index, + "model_name": config.model.name, + "best_epoch": trained.best_epoch, + "state_dict": trained.best_state, + }, + checkpoint_dir / f"fold{fold.index}.pt", + ) + + eval_started = time.monotonic() + use_amp = config.train.amp and device.type == "cuda" + logits, _ = predict_logits(model, held_loader, device, amp=use_amp) + probs = torch.softmax(torch.from_numpy(logits), dim=1).numpy() + eval_seconds = time.monotonic() - eval_started + held_acc = float(np.mean(probs.argmax(axis=1) == labels[held_idx])) + logger.info( + "%s: 最良エポック %d（val_macro_f1 %.4f）・held_out accuracy %.4f" + "（学習 %.1f 秒・評価 %.1f 秒）", + prefix, + trained.best_epoch, + trained.best_val_macro_f1, + held_acc, + trained.seconds, + eval_seconds, + ) + + result.folds.append( + FoldResult( + index=fold.index, + seed=seed, + train_patients=train_patients, + val_patients=val_patients, + held_out_patients=fold.held_out, + group_counts=group_counts, + history=trained.history, + best_epoch=trained.best_epoch, + best_val_macro_f1=trained.best_val_macro_f1, + train_seconds=trained.seconds, + eval_seconds=eval_seconds, + held_out_indices=held_idx, + held_out_probs=probs, + ) + ) + del model, trained + if device.type == "cuda": + torch.cuda.empty_cache() + return result + + +def _make_loader( + images: np.ndarray, + labels: np.ndarray, + config: TrainConfig, + device: torch.device, + *, + train: bool, + seed: int, +) -> DataLoader: + dataset = TongueImageDataset(images, labels, augment=train) + generator = torch.Generator() + generator.manual_seed(seed) + return DataLoader( + dataset, + batch_size=config.batch_size, + shuffle=train, + # BatchNorm を含むモデルで最後のバッチが 1 件になると学習が止まるため，学習時は端数を捨てる + # （件数がバッチサイズ以下のときは捨てない）． + drop_last=train and len(dataset) > config.batch_size, + num_workers=config.num_workers, + pin_memory=device.type == "cuda", + persistent_workers=config.num_workers > 0, + generator=generator, + ) diff --git a/src/tias_judge_model/training/loop.py b/src/tias_judge_model/training/loop.py new file mode 100644 index 0000000..d67f706 --- /dev/null +++ b/src/tias_judge_model/training/loop.py @@ -0,0 +1,279 @@ +"""1 つの学習（1 fold 分）のループと推論（SPEC_02「モデル」「学習の設定」）. + +AdamW ＋ 線形 warmup → cosine 減衰（ステップ単位）．CUDA のときだけ AMP を使う． +エポックごとに検証の macro F1 を測り，最良エポックの重みを保持し，`patience` エポック改善が +無ければ打ち切る．ログには集計値（損失・F1・時間）だけを出す． +""" + +from __future__ import annotations + +import logging +import math +import random +import time +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader + +from tias_judge_model.data.dataset import normalize_batch +from tias_judge_model.evaluation.metrics import macro_f1 +from tias_judge_model.schema import TrainConfig + +logger = logging.getLogger(__name__) + + +@dataclass +class EpochRecord: + """1 エポックの記録（`history.csv` の 1 行）. + + Attributes: + epoch: エポック番号（1 始まり）． + train_loss: 学習損失の平均． + val_loss: 検証損失の平均． + val_macro_f1: 検証の macro F1． + lr: エポック終了時の学習率． + seconds: このエポックの所要秒数（学習 ＋ 検証）． + """ + + epoch: int + train_loss: float + val_loss: float + val_macro_f1: float + lr: float + seconds: float + + +@dataclass +class TrainResult: + """1 回の学習の結果. + + Attributes: + best_state: 最良エポックの重み（CPU 上の state_dict）． + best_epoch: 最良エポック番号（1 始まり）． + best_val_macro_f1: 最良エポックの検証 macro F1． + history: エポックごとの記録． + seconds: 学習全体の所要秒数． + """ + + best_state: dict[str, torch.Tensor] + best_epoch: int + best_val_macro_f1: float + history: list[EpochRecord] = field(default_factory=list) + seconds: float = 0.0 + + @property + def epochs_run(self) -> int: + """実際に回したエポック数.""" + return len(self.history) + + +def seed_everything(seed: int) -> None: + """`random`・`numpy`・`torch`（CUDA を含む）の乱数 seed を固定する. + + Args: + seed: 乱数 seed． + """ + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def select_device() -> torch.device: + """CUDA があれば `cuda`（`CUDA_VISIBLE_DEVICES` で割り当てられた GPU），無ければ CPU を返す.""" + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def device_name(device: torch.device) -> str: + """デバイスの表示名（GPU 名または `cpu`）を返す.""" + if device.type == "cuda": + return torch.cuda.get_device_name(device) + return "cpu" + + +def make_scheduler( + optimizer: torch.optim.Optimizer, *, warmup_steps: int, total_steps: int +) -> torch.optim.lr_scheduler.LambdaLR: + """線形 warmup の後に cosine で 0 へ減衰する学習率スケジューラ（ステップ単位）を作る. + + Args: + optimizer: 最適化器． + warmup_steps: warmup のステップ数． + total_steps: 全ステップ数． + + Returns: + ステップごとに `step()` を呼ぶスケジューラ． + """ + + def factor(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return (step + 1) / warmup_steps + progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) + return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, factor) + + +def train_model( + model: nn.Module, + train_loader: DataLoader, + val_loader: DataLoader, + config: TrainConfig, + device: torch.device, + *, + log_prefix: str = "", + on_epoch_end: Callable[[EpochRecord], None] | None = None, +) -> TrainResult: + """モデルを学習し，検証 macro F1 が最良のエポックの重みを返す（早期終了あり）. + + Args: + model: 学習するモデル（`device` 上に置いておく）． + train_loader: 学習用（`(B, 3, H, W)` の uint8 画像と `(B,)` のクラス ID を返す）． + val_loader: 検証用（同上）． + config: 学習の設定． + device: デバイス． + log_prefix: ログの先頭に付ける文字列（例: `fold 1/5`）． + on_epoch_end: エポックごとに呼ぶ関数（任意）． + + Returns: + 学習の結果．`model` の重みは最良エポックのものに戻しておく． + """ + started = time.monotonic() + use_amp = config.amp and device.type == "cuda" + optimizer = torch.optim.AdamW( + model.parameters(), lr=config.lr, weight_decay=config.weight_decay + ) + steps_per_epoch = max(1, len(train_loader)) + scheduler = make_scheduler( + optimizer, + warmup_steps=config.warmup_epochs * steps_per_epoch, + total_steps=config.max_epochs * steps_per_epoch, + ) + scaler = torch.amp.GradScaler("cuda", enabled=use_amp) + criterion = nn.CrossEntropyLoss() + + best_state: dict[str, torch.Tensor] | None = None + best_epoch = 0 + best_f1 = -1.0 + history: list[EpochRecord] = [] + for epoch in range(1, config.max_epochs + 1): + epoch_started = time.monotonic() + model.train() + loss_sum = 0.0 + count = 0 + for images, labels in train_loader: + inputs = normalize_batch(images.to(device, non_blocking=True)) + targets = labels.to(device, non_blocking=True) + optimizer.zero_grad(set_to_none=True) + with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp): + loss = criterion(model(inputs), targets) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + scheduler.step() + loss_sum += float(loss.detach()) * len(targets) + count += len(targets) + + val_logits, val_labels = predict_logits(model, val_loader, device, amp=use_amp) + val_loss = float( + criterion(torch.from_numpy(val_logits), torch.from_numpy(val_labels)).item() + ) + val_f1 = macro_f1(val_labels, val_logits.argmax(axis=1)) + record = EpochRecord( + epoch=epoch, + train_loss=loss_sum / max(1, count), + val_loss=val_loss, + val_macro_f1=val_f1, + lr=float(optimizer.param_groups[0]["lr"]), + seconds=time.monotonic() - epoch_started, + ) + history.append(record) + if on_epoch_end is not None: + on_epoch_end(record) + + improved = val_f1 > best_f1 + if improved: + best_f1 = val_f1 + best_epoch = epoch + best_state = {k: v.detach().to("cpu", copy=True) for k, v in model.state_dict().items()} + logger.info( + "%s epoch %d/%d: train_loss %.4f・val_loss %.4f・val_macro_f1 %.4f・lr %.2e" + "（%.1f 秒・経過 %.1f 秒）%s", + log_prefix, + epoch, + config.max_epochs, + record.train_loss, + record.val_loss, + record.val_macro_f1, + record.lr, + record.seconds, + time.monotonic() - started, + " *best" if improved else "", + ) + if epoch - best_epoch >= config.patience: + logger.info( + "%s 検証 macro F1 が %d エポック改善しないため打ち切る", log_prefix, config.patience + ) + break + + if best_state is None: + raise RuntimeError("学習が 1 エポックも行われなかった") + model.load_state_dict(best_state) + return TrainResult( + best_state=best_state, + best_epoch=best_epoch, + best_val_macro_f1=best_f1, + history=history, + seconds=time.monotonic() - started, + ) + + +@torch.no_grad() +def predict_logits( + model: nn.Module, loader: DataLoader, device: torch.device, *, amp: bool +) -> tuple[np.ndarray, np.ndarray]: + """データ全体のロジットと正解を返す（推論モード）. + + Args: + model: モデル． + loader: `(B, 3, H, W)` の uint8 画像と `(B,)` のクラス ID を返す DataLoader + （shuffle なし）． + device: デバイス． + amp: 混合精度を使うか（CUDA のときだけ有効にすること）． + + Returns: + (`(N, C)` の float32 ロジット, `(N,)` の int64 正解クラス ID)．並びは loader の順． + """ + model.eval() + logits: list[np.ndarray] = [] + labels: list[np.ndarray] = [] + for images, targets in loader: + inputs = normalize_batch(images.to(device, non_blocking=True)) + with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=amp): + out = model(inputs) + logits.append(out.float().cpu().numpy()) + labels.append(targets.numpy().astype(np.int64)) + if not logits: + return np.zeros((0, 0), dtype=np.float32), np.zeros((0,), dtype=np.int64) + return np.concatenate(logits), np.concatenate(labels) + + +def summarize_history(history: list[EpochRecord]) -> dict[str, Any]: + """エポックの記録から所要時間の要約（エポック数・1 エポックあたり平均秒数）を返す. + + Args: + history: エポックごとの記録． + + Returns: + `{"epochs_run", "mean_epoch_seconds"}`． + """ + seconds = [h.seconds for h in history] + return { + "epochs_run": len(history), + "mean_epoch_seconds": float(np.mean(seconds)) if seconds else None, + }