mask_equivalence.py

#!/usr/bin/env python3
"""bool のマスクと等価なのは -inf の float マスクであって、1.0 のそれではない。

モデルもチェックポイントも要らない。PyTorch だけで完結する。

    pip install torch
    python mask_equivalence.py

`train/README.md` は ONNX 変換のために parseq のマスクを
`dtype=torch.bool` から `dtype=torch.float` へ変えるよう指示している。
変更自体は必要だが、`torch.ones` のまま float にすると意味が変わる。

PyTorch の説明（`torch.nn.MultiheadAttention` の docstring）:

    Binary and float masks are supported. For a binary mask, a ``True`` value
    indicates that the corresponding position is not allowed to attend. For a
    float mask, the mask values will be added to the attention weight.

そして実装（`torch/nn/functional.py` の `_canonical_mask`）:

    if not _mask_is_float:
        mask = torch.zeros_like(mask, dtype=target_type).masked_fill_(
            mask, float("-inf")
        )

つまり bool の True は -inf に変換される。1.0 ではない。
"""

import torch
import torch.nn as nn


def main() -> None:
    torch.manual_seed(0)
    n, dim, heads = 6, 8, 2
    attention = nn.MultiheadAttention(dim, heads, batch_first=True).eval()
    query, key_value = torch.randn(1, n, dim), torch.randn(1, n, dim)

    upper = torch.triu(torch.ones(n, n, dtype=torch.bool), 1)
    masks = {
        "bool（parseq 本来）": upper,
        "float ones（README の指示）": torch.triu(torch.ones(n, n), 1),
        "float -inf（提案）": torch.triu(torch.full((n, n), float("-inf")), 1),
    }

    print("== 1. PyTorch が bool マスクを何に変換するか ==")
    converted = torch.zeros_like(upper, dtype=torch.float32).masked_fill_(upper, float("-inf"))
    print("   bool マスク:")
    print(upper.int())
    print("   PyTorch が内部で作る float マスク:")
    print(converted)
    print(f"   torch.triu(torch.full(...,-inf),1) と同一か: "
          f"{torch.equal(converted, masks['float -inf（提案）'])}")

    print("\n== 2. それぞれのマスクで注意の出力を比べる ==")
    outputs = {}
    for label, mask in masks.items():
        with torch.inference_mode():
            outputs[label] = attention(query, key_value, key_value, attn_mask=mask)[0]
    base = outputs["bool（parseq 本来）"]
    for label, out in outputs.items():
        print(f"   {label:26s} bool との max_abs {(out - base).abs().max().item():.3e}")

    print("\n== 3. cloze マスク（parseq の refine 段が使う形）でも同じか ==")
    # 位置 i の問い合わせから、答えの入っている位置 i+1 だけを隠す。
    # 上三角のうち 2 つ以上先を開けると、対角のすぐ上だけが残る。
    def cloze(mask, open_value):
        mask = mask.clone()
        mask[torch.triu(torch.ones(n, n, dtype=torch.bool), 2)] = open_value
        return mask

    cloze_masks = {
        "bool（parseq 本来）": cloze(upper, False),
        "float ones（README の指示）": cloze(masks["float ones（README の指示）"], 0.0),
        "float -inf（提案）": cloze(masks["float -inf（提案）"], 0.0),
    }
    print("   float -inf 版の cloze マスク（-inf が隠す位置）:")
    print(cloze_masks["float -inf（提案）"])
    print("   float ones 版（同じ位置に +1 が入る = 隠すどころか加点する）:")
    print(cloze_masks["float ones（README の指示）"])

    outputs = {}
    for label, mask in cloze_masks.items():
        with torch.inference_mode():
            outputs[label] = attention(query, key_value, key_value, attn_mask=mask)[0]
    base = outputs["bool（parseq 本来）"]
    print()
    for label, out in outputs.items():
        print(f"   {label:26s} bool との max_abs {(out - base).abs().max().item():.3e}")

    print("\n== 4. 注意重みそのものを見る ==")
    weights = {}
    for label, mask in cloze_masks.items():
        with torch.inference_mode():
            weights[label] = attention(query, key_value, key_value, attn_mask=mask)[1]
    print("   位置 0 の問い合わせが、各位置にどれだけ注意を向けるか（隠すべきは位置 1）:")
    for label, w in weights.items():
        row = w[0, 0].tolist()
        marks = " ".join(f"{v:.3f}" for v in row)
        hidden = row[1]
        verdict = "隠せている" if hidden < 1e-6 else (
            "行内で最大に注意している" if hidden == max(row) else "隠せていない")
        print(f"     {label:26s} [{marks}]  位置 1 = {hidden:.3f} -> {verdict}")


if __name__ == "__main__":
    main()

onnx_weights.py

#!/usr/bin/env python3
"""配布 ONNX 認識器と公開チェックポイントが別の重みであることの再現。

ndl-lab/ndlocr-lite への報告に添える単体スクリプト。このリポジトリの他の
ファイルには依存しない。

  pip install torch onnx onnxruntime pyyaml pillow
  git clone https://github.com/baudm/parseq   # 変更を加えないもの

  python repro_weights.py \
    --parseq ./parseq \
    --onnx  parseq-ndl-24x256-30-tiny-189epoch-tegaki3-r8data-202604.onnx \
    --checkpoint parseq-ndl-24x256-30-tiny-189epoch-tegaki3-r8data-202604.ckpt \
    --charset NDLmoji.yaml

やっていること:

1. 配布 ONNX の初期化子を parseq の PyTorch モデルに読み戻す
2. 読み戻した模型が配布 ONNX と同じ読みをすることを確かめる（測定系の対照）
3. その重みを公開チェックポイントと比べる（state_dict と SWA 平均の両方）
4. --lines-dir があれば CER も測る

チェックポイントは pickle 形式なので、信頼できる公式ファイルにのみ使うこと。
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import torch
import yaml
from onnx import numpy_helper


def load_system(parseq: Path, checkpoint: Path, charset: str):
    sys.path.insert(0, str(parseq))
    from strhub.models.utils import load_from_checkpoint

    original = torch.load

    def trusted(*args, **kwargs):
        kwargs["weights_only"] = False
        return original(*args, **kwargs)

    torch.load = trusted
    try:
        return load_from_checkpoint(str(checkpoint), charset_test=charset).eval()
    finally:
        torch.load = original


def module_path(node_name: str) -> str:
    """ノード名からモジュールのパスを作る。

    '/encoder/blocks/blocks.0/attn/qkv/MatMul' -> 'encoder.blocks.0.attn.qkv'

    書き出しは子の添字を「親の名前 + 添字」という区切りで書くので、親が
    二度現れる。片方を落とす。
    """
    parts = [p for p in node_name.split("/") if p][:-1]
    out: list[str] = []
    for part in parts:
        if out and part.startswith(out[-1] + "."):
            out[-1] = part
        else:
            out.append(part)
    return ".".join(out)


def transplant(system, onnx_path: Path) -> tuple[dict, list[str]]:
    """配布 ONNX の重みを state_dict に読み戻す。"""
    graph = onnx.load(str(onnx_path)).graph
    init = {i.name: numpy_helper.to_array(i) for i in graph.initializer}
    state = system.state_dict()
    keys = set(state)
    new: dict[str, torch.Tensor] = {}

    # (1) 名前がそのまま残っているもの。
    for key in state:
        if key in init and tuple(init[key].shape) == tuple(state[key].shape):
            new[key] = torch.from_numpy(init[key].copy())

    # ノードごとに、その定数入力を集めておく。
    # 定数は MatMul では入力 1、Add では入力 0 にある。
    scoped: dict[str, dict] = {}
    for node in graph.node:
        for name in node.input:
            if name in init:
                scope, _, op = node.name.rpartition("/")
                scoped.setdefault(scope, {})[op] = init[name]
                break

    # (2) MultiheadAttention の in_proj。parseq はクエリと key/value に別の
    #     テンソルを渡すので、書き出しは 3 分割ではなく「クエリ」と
    #     「key/value」の 2 本に割る。その順に連結して戻す。
    fused: set[str] = set()
    for scope, ops in scoped.items():
        if not (scope.endswith("self_attn") or scope.endswith("cross_attn")):
            continue
        key = "model." + module_path(scope + "/x") + ".in_proj_weight"
        query, key_value = ops.get("MatMul"), ops.get("MatMul_1")
        if key not in keys or query is None or key_value is None:
            continue
        new[key] = torch.from_numpy(np.concatenate([query.T, key_value.T], axis=0).copy())
        bias_key = key.replace("in_proj_weight", "in_proj_bias")
        bq, bkv = ops.get("Add"), ops.get("Add_1")
        if bias_key in keys and bq is not None and bkv is not None:
            new[bias_key] = torch.from_numpy(np.concatenate([bq, bkv], axis=0).copy())
        fused.add(scope)

    # (3) nn.Linear は転置された無名の MatMul 定数になっている。
    for node in graph.node:
        if node.op_type == "MatMul" and node.input[1] in init:
            if node.name.rpartition("/")[0] in fused:
                continue
            key = "model." + module_path(node.name) + ".weight"
            if key in keys and key not in new:
                new[key] = torch.from_numpy(init[node.input[1]].T.copy())

    # (4) pos_queries は展開された復号ループが切り出す元の定数として残る。
    #     形が合うだけで決めず、各ステップのスライスがその前置になることを確かめる。
    key = "model.pos_queries"
    if key in keys and key not in new:
        want = tuple(state[key].shape)
        for name, array in init.items():
            if tuple(array.shape) != want or name in new:
                continue
            slices = [
                a for a in init.values()
                if a.ndim == 3 and a.shape[0] == 1 and a.shape[2] == want[2] and 0 < a.shape[1] <= want[1]
            ]
            if slices and all(np.array_equal(a, array[:, : a.shape[1]]) for a in slices):
                new[key] = torch.from_numpy(array.copy())
                break

    system.load_state_dict(new, strict=False)
    return new, sorted(keys - set(new))


_ORIGINAL_FORWARD_STREAM = None


def set_mask(parseq: Path, shipped: bool) -> None:
    """マスクの扱いを、配布 ONNX と同じにするか parseq のままにするか切り替える。

    parseq の attn_mask は bool で、PyTorch はこれを「禁止」と読んで -inf を
    足す。train/README.md は ONNX 変換のために dtype=torch.float へ変える
    よう指示しており、float のマスクは加算されるので 1.0 が加点になる。
    配布 ONNX を再現するにはそちらに合わせる必要があるが、公開チェックポイント
    を評価するときは parseq のままでなければならないので、模型ごとに切り替える。
    """
    global _ORIGINAL_FORWARD_STREAM
    sys.path.insert(0, str(parseq))
    from strhub.models.parseq import modules

    if _ORIGINAL_FORWARD_STREAM is None:
        _ORIGINAL_FORWARD_STREAM = modules.DecoderLayer.forward_stream
    original = _ORIGINAL_FORWARD_STREAM

    def forward_stream(self, tgt, tgt_norm, tgt_kv, memory, tgt_mask, pad):
        if tgt_mask is not None and tgt_mask.dtype == torch.bool:
            tgt_mask = tgt_mask.to(torch.float32)
        return original(self, tgt, tgt_norm, tgt_kv, memory, tgt_mask, pad)

    modules.DecoderLayer.forward_stream = forward_stream if shipped else original


def read(model, x: np.ndarray, parseq: Path, shipped: bool) -> np.ndarray:
    set_mask(parseq, shipped)
    with torch.inference_mode():
        return model(torch.from_numpy(x), max_length=model.model.max_label_length).numpy().argmax(-1)[0]


def decode(ids, charset: str) -> str:
    out = []
    for token in ids:
        if token == 0:  # EOS
            break
        out.append(charset[token - 1])
    return "".join(out)


def edit_distance(a: str, b: str) -> int:
    previous = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        current = [i]
        for j, cb in enumerate(b, 1):
            current.append(min(previous[j] + 1, current[j - 1] + 1, previous[j - 1] + (ca != cb)))
        previous = current
    return previous[-1]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--parseq", type=Path, required=True, help="baudm/parseq の clone（無改変）")
    parser.add_argument("--onnx", type=Path, required=True)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--charset", type=Path, required=True, help="NDLmoji.yaml")
    parser.add_argument("--samples", type=int, default=20, help="対照に使う乱数入力の数")
    parser.add_argument(
        "--lines-dir",
        type=Path,
        help="CER も測る場合。manifest.json と <name>.png を含むディレクトリ "
        "(ndl-lab/ndl-minhon-ocrdataset から作れる)",
    )
    args = parser.parse_args()

    charset = yaml.safe_load(args.charset.read_text())["model"]["charset_test"]
    session = ort.InferenceSession(str(args.onnx), providers=["CPUExecutionProvider"])
    in_name = session.get_inputs()[0].name
    _, _, height, width = session.get_inputs()[0].shape

    print("== 1. 配布 ONNX の重みを PyTorch に読み戻す ==")
    moved = load_system(args.parseq, args.checkpoint, charset)
    filled, missing = transplant(moved, args.onnx)
    total = len(moved.state_dict())
    print(f"   {len(filled)}/{total} 個を移植、未充填 {len(missing)} 個")
    for key in missing:
        print(f"     未充填: {key}")

    print("\n== 2. 対照: 読み戻した模型は配布 ONNX と同じ読みをするか ==")
    rng = np.random.default_rng(0)
    agree = 0
    for _ in range(args.samples):
        x = rng.standard_normal((1, 3, height, width), dtype=np.float32)
        want = session.run(None, {in_name: x})[0].argmax(-1)[0]
        got = read(moved, x, args.parseq, shipped=True)
        agree += bool(np.array_equal(want, got))
    print(f"   {agree}/{args.samples} 件で一致")
    print("   (一致すれば、ONNX 側の重みは欠けなく読み出せている)")

    print("\n== 3. その重みを公開チェックポイントと比べる ==")
    reference = load_system(args.parseq, args.checkpoint, charset)
    moved_state = moved.state_dict()

    # チェックポイントは重みを 2 通り持っている。load_from_checkpoint が読む
    # state_dict と、StochasticWeightAveraging が別に保持する平均である。
    # 書き出し元がどちらかは分からないので、両方と比べる。
    raw = torch.load(str(args.checkpoint), map_location="cpu", weights_only=False)
    print(f"   チェックポイントの epoch: {raw.get('epoch', '不明')}")
    stem = args.checkpoint.stem
    if (match := __import__("re").search(r"-(\d+)epoch", stem)) and raw.get("epoch") is not None:
        named, actual = int(match.group(1)), raw["epoch"]
        if named != actual:
            print(f"   ※ ファイル名は {named} epoch ですが、中身は {actual} epoch です")

    candidates = {"state_dict": reference.state_dict()}
    swa = raw.get("callbacks", {}).get("StochasticWeightAveraging", {})
    if "average_model_state" in swa:
        candidates["SWA 平均"] = swa["average_model_state"]

    for label, published in candidates.items():
        keys = [k for k in moved_state if k in published]
        exact = sum(1 for k in keys if torch.equal(published[k].float(), moved_state[k].float()))
        worst = max((published[k].float() - moved_state[k].float()).abs().max().item() for k in keys)
        cosines = []
        for k in keys:
            a = moved_state[k].float().flatten().double()
            b = published[k].float().flatten().double()
            if a.norm() > 1e-12 and b.norm() > 1e-12:
                cosines.append((a @ b / (a.norm() * b.norm())).item())
        verdict = "一致" if exact == len(keys) else "別物"
        print(f"   vs {label:11s}: 完全一致 {exact:3d}/{len(keys)} | 最大差 {worst:.2e}"
              f" | cosine 中央 {np.median(cosines):.5f}  -> {verdict}")

    if not args.lines_dir:
        print("\n(CER を測るには --lines-dir を渡してください)")
        return

    print("\n== 4. CER ==")
    from PIL import Image

    rows = [
        r for r in json.loads((args.lines_dir / "manifest.json").read_text(encoding="utf-8"))
        if 0 < len(r["text"]) <= moved.model.max_label_length and set(r["text"]) <= set(charset)
    ]
    print(f"   {len(rows)} 行を採点", flush=True)

    def prep(path: Path) -> np.ndarray:
        image = Image.open(path).convert("RGB")
        if image.height > image.width:  # 縦長の切り出しは回す
            image = image.transpose(Image.ROTATE_90)
        a = np.asarray(image.resize((width, height), Image.Resampling.BILINEAR), dtype=np.float32)
        a = a[:, :, ::-1].copy()  # 製品側は BGR で与えている
        return np.transpose(a / 127.5 - 1.0, (2, 0, 1))[None].astype(np.float32)

    # 読み戻した模型は配布 ONNX のマスクで、公開チェックポイントは parseq の
    # ままのマスクで動かす。それぞれ本来の姿で評価するため。
    arms = {"配布 ONNX": (None, None), "読み戻したもの(対照)": (moved, True), "公開チェックポイント": (reference, False)}
    totals = {name: [0, 0] for name in arms}
    for index, row in enumerate(rows, 1):
        path = args.lines_dir / f"{row['name']}.png"
        if not path.exists():
            continue
        x = prep(path)
        for name, (model, shipped) in arms.items():
            if model is None:
                ids = session.run(None, {in_name: x})[0].argmax(-1)[0]
            else:
                ids = read(model, x, args.parseq, shipped)
            totals[name][0] += edit_distance(decode(ids, charset), row["text"])
            totals[name][1] += len(row["text"])
        if index % 500 == 0:
            print(f"   {index}/{len(rows)} " + " ".join(f"{n} {t[0]/max(t[1],1):.4f}" for n, t in totals.items()), flush=True)
    print()
    for name, (errors, length) in totals.items():
        print(f"   {name:22s} CER {errors/max(length,1):.4f}")


if __name__ == "__main__":
    main()