Archive standalone checkpoint converter

convert_checkpoint.py

@@ -0,0 +1,79 @@
+"""Convert an official ST-EEGFormer MAE checkpoint to braindecode format.
+
+The braindecode ``STEEGFormer`` loads braindecode-format state dicts directly,
+so this one-off converter is what produced the re-hosted ``model.safetensors``
+in this repo. It remaps the upstream ``timm`` keys (drops the MAE decoder and
+downstream-only keys, renames the blocks, and splits the fused ``attn.qkv``
+into the separate queries/keys/values of braindecode's ``MultiHeadAttention``).
+
+Usage:
+    python convert_checkpoint.py checkpoint-300.pth ./out \\
+        --embed-dim 512 --depth 8 --num-heads 8 --n-chans-pos 145
+"""
+import argparse
+from collections import OrderedDict
+
+import torch
+
+from braindecode.models import STEEGFormer
+
+_BLOCK = {
+    "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_PREFIX = ("decoder_", "dec_", "mask_token", "enc_temporal_emd")
+_DROP_EXACT = {"pos_embed", "fc_norm.weight", "fc_norm.bias", "head.weight", "head.bias"}
+
+
+def remap(state_dict, embed_dim):
+    if isinstance(state_dict.get("model"), dict):
+        state_dict = state_dict["model"]
+    e, out = embed_dim, OrderedDict()
+    for k, v in state_dict.items():
+        if k.startswith(_DROP_PREFIX) or k in _DROP_EXACT:
+            continue
+        if k == "enc_channel_emd.channel_transformation.weight":
+            out["channel_pos.embedding.weight"] = v
+        elif k.startswith("blocks."):
+            _, i, rest = k.split(".", 2)
+            d = f"encoder.{i}."
+            if rest in ("attn.qkv.weight", "attn.qkv.bias"):
+                s = rest.rsplit(".", 1)[1]
+                out[f"{d}0.fn.1.queries.{s}"] = v[:e]
+                out[f"{d}0.fn.1.keys.{s}"] = v[e : 2 * e]
+                out[f"{d}0.fn.1.values.{s}"] = v[2 * e :]
+            else:
+                for a, b in _BLOCK.items():
+                    if rest.startswith(a + "."):
+                        out[d + b + rest[len(a) :]] = v
+                        break
+        else:  # cls_token, patch_embed.proj.*, norm.* keep their names
+            out[k] = v
+    return out
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("checkpoint")
+    p.add_argument("out_dir")
+    p.add_argument("--embed-dim", type=int, required=True)
+    p.add_argument("--depth", type=int, required=True)
+    p.add_argument("--num-heads", type=int, required=True)
+    p.add_argument("--n-chans-pos", type=int, default=145)
+    a = p.parse_args()
+    ck = torch.load(a.checkpoint, map_location="cpu", weights_only=False)
+    model = STEEGFormer(
+        n_chans=22, n_outputs=4, n_times=1000, embed_dim=a.embed_dim,
+        depth=a.depth, num_heads=a.num_heads, n_chans_pos=a.n_chans_pos,
+    )
+    res = model.load_state_dict(remap(ck, a.embed_dim), strict=False)
+    assert not res.unexpected_keys, f"unexpected: {res.unexpected_keys[:5]}"
+    model.save_pretrained(a.out_dir)
+    print(f"saved braindecode-format checkpoint to {a.out_dir}")
+
+
+if __name__ == "__main__":
+    main()