Add ST-EEGFormer conversion script

convert_steegformer_checkpoints.py

@@ -0,0 +1,228 @@
+"""Convert and re-host ST-EEGFormer checkpoints on the Hugging Face Hub.
+
+The official ST-EEGFormer checkpoints are MAE pre-training checkpoints from
+``LiuyinYang1101/STEEGFormer``. They are stored under timm-style encoder keys,
+include decoder-only tensors for the pre-training objective, and fuse attention
+queries/keys/values in ``attn.qkv``. This script converts those files into the
+braindecode model format used by :class:`~braindecode.models.STEEGFormer`.
+
+The conversion is intentionally kept out of ``STEEGFormer.load_state_dict``:
+runtime loading expects braindecode-format checkpoints, while this one-off
+remapping script is archived with the re-hosted weights.
+
+Run from the repo root with the official checkpoint available locally::
+
+    python hf_assets/model_cards/convert_steegformer_checkpoints.py \
+        --src STEEGFormer_small.pth --variant small
+
+Add ``--push`` (requires write access to the ``braindecode`` HF org) to upload
+the converted folder, including this script, to the corresponding Hub repo.
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+import sys
+from collections import OrderedDict
+from pathlib import Path
+from typing import Any
+
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+if (_REPO_ROOT / "braindecode" / "models").exists():
+    sys.path.insert(0, str(_REPO_ROOT))
+
+VARIANTS: dict[str, dict[str, Any]] = {
+    "small": {
+        "embed_dim": 512,
+        "depth": 8,
+        "num_heads": 8,
+        "n_chans_pos": 145,
+    },
+    "base": {
+        "embed_dim": 768,
+        "depth": 12,
+        "num_heads": 12,
+        "n_chans_pos": 145,
+    },
+    "large": {
+        "embed_dim": 1024,
+        "depth": 24,
+        "num_heads": 16,
+        "n_chans_pos": 145,
+    },
+    "largeV2": {
+        "embed_dim": 1024,
+        "depth": 24,
+        "num_heads": 16,
+        "n_chans_pos": 256,
+    },
+}
+
+HUB_REPOS = {
+    "small": "braindecode/STEEGFormer-small",
+    "base": "braindecode/STEEGFormer-base",
+    "large": "braindecode/STEEGFormer-large",
+    "largeV2": "braindecode/STEEGFormer-largeV2",
+}
+
+# timm key inside one block -> this module's encoder-block key. The fused
+# ``attn.qkv`` tensor is split separately.
+_TIMM_BLOCK_RENAMES = {
+    "norm1": "0.fn.0",
+    "attn.proj": "0.fn.1.projection",
+    "norm2": "1.fn.0",
+    "mlp.fc1": "1.fn.1.0",
+    "mlp.fc2": "1.fn.1.3",
+}
+
+_DROP_PREFIXES = ("decoder_", "dec_", "mask_token", "enc_temporal_emd")
+_DROP_EXACT = frozenset(
+    {"pos_embed", "fc_norm.weight", "fc_norm.bias", "head.weight", "head.bias"}
+)
+_EXPECTED_MISSING = {"final_layer.weight", "final_layer.bias"}
+
+
+def _load_checkpoint(src: Path) -> dict[str, Any]:
+    try:
+        return torch.load(src, map_location="cpu", weights_only=True)
+    except Exception:
+        return torch.load(src, map_location="cpu", weights_only=False)
+
+
+def _unwrap_state_dict(checkpoint: dict[str, Any]) -> dict[str, torch.Tensor]:
+    for key in ("model", "state_dict"):
+        nested = checkpoint.get(key)
+        if isinstance(nested, dict):
+            return nested
+    return checkpoint
+
+
+def remap_official_state_dict(
+    checkpoint: dict[str, Any], embed_dim: int
+) -> OrderedDict[str, torch.Tensor]:
+    """Return a braindecode-format state dict from an official timm checkpoint."""
+    state_dict = _unwrap_state_dict(checkpoint)
+    if not any("attn.qkv" in key for key in state_dict):
+        raise ValueError("Expected an official timm checkpoint with attn.qkv keys.")
+
+    remapped: OrderedDict[str, torch.Tensor] = OrderedDict()
+    unknown_block_keys: list[str] = []
+
+    for key, value in state_dict.items():
+        if key.startswith(_DROP_PREFIXES) or key in _DROP_EXACT:
+            continue
+
+        if key == "enc_channel_emd.channel_transformation.weight":
+            remapped["channel_pos.embedding.weight"] = value
+        elif key.startswith("blocks."):
+            _, idx, rest = key.split(".", 2)
+            dst = f"encoder.{idx}."
+            if rest in ("attn.qkv.weight", "attn.qkv.bias"):
+                if value.shape[0] != 3 * embed_dim:
+                    raise ValueError(
+                        f"{key} has first dimension {value.shape[0]}, expected "
+                        f"{3 * embed_dim} for embed_dim={embed_dim}."
+                    )
+                suffix = rest.rsplit(".", 1)[1]
+                remapped[f"{dst}0.fn.1.queries.{suffix}"] = value[:embed_dim]
+                remapped[f"{dst}0.fn.1.keys.{suffix}"] = value[
+                    embed_dim : 2 * embed_dim
+                ]
+                remapped[f"{dst}0.fn.1.values.{suffix}"] = value[2 * embed_dim :]
+            else:
+                for src, new in _TIMM_BLOCK_RENAMES.items():
+                    if rest.startswith(src + "."):
+                        remapped[dst + new + rest[len(src) :]] = value
+                        break
+                else:
+                    unknown_block_keys.append(key)
+        else:
+            # cls_token, patch_embed.proj.*, and norm.* share braindecode names.
+            remapped[key] = value
+
+    if unknown_block_keys:
+        shown = ", ".join(unknown_block_keys[:8])
+        if len(unknown_block_keys) > 8:
+            shown += ", ..."
+        raise RuntimeError(f"Unmapped timm block keys: {shown}")
+
+    return remapped
+
+
+def convert(src: Path, variant: str, out: Path, n_outputs: int = 4) -> None:
+    """Convert one official checkpoint into a braindecode pretrained folder."""
+    from braindecode.models import STEEGFormer
+
+    config = dict(VARIANTS[variant])
+    checkpoint = _load_checkpoint(src)
+    state_dict = remap_official_state_dict(checkpoint, embed_dim=config["embed_dim"])
+
+    model = STEEGFormer(
+        n_chans=22,
+        n_times=1000,
+        n_outputs=n_outputs,
+        patch_size=16,
+        mlp_ratio=4,
+        **config,
+    )
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+    extra_missing = sorted(set(missing) - _EXPECTED_MISSING)
+    if extra_missing or unexpected:
+        raise RuntimeError(
+            f"Converted checkpoint mismatch: missing={extra_missing}, "
+            f"unexpected={list(unexpected)}"
+        )
+
+    out.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out)
+    shutil.copy2(Path(__file__).resolve(), out / Path(__file__).name)
+    print(f"Saved {variant} to {out} ({len(state_dict)} tensors loaded)")
+
+
+def push(out: Path, variant: str) -> None:
+    """Upload a converted folder to the matching braindecode Hub repo."""
+    from huggingface_hub import HfApi
+
+    repo_id = HUB_REPOS[variant]
+    api = HfApi()
+    api.create_repo(repo_id, repo_type="model", exist_ok=True)
+    api.upload_folder(
+        repo_id=repo_id,
+        folder_path=str(out),
+        commit_message="Add ST-EEGFormer braindecode checkpoint conversion",
+    )
+    print(f"Pushed {out} -> {repo_id}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--src", type=Path, required=True)
+    parser.add_argument("--variant", choices=sorted(VARIANTS), required=True)
+    parser.add_argument(
+        "--out",
+        type=Path,
+        default=None,
+        help="output folder (default: hf_export/STEEGFormer-<variant>)",
+    )
+    parser.add_argument("--n-outputs", type=int, default=4)
+    parser.add_argument(
+        "--push",
+        action="store_true",
+        help="upload to the matching braindecode/STEEGFormer-* repo",
+    )
+    args = parser.parse_args()
+
+    out = args.out or Path("hf_export") / f"STEEGFormer-{args.variant}"
+    convert(args.src, args.variant, out, n_outputs=args.n_outputs)
+    if args.push:
+        push(out, args.variant)
+
+
+if __name__ == "__main__":
+    main()