Remove duplicate converter (keep convert_steegformer_checkpoints.py)

convert_checkpoint.py

@@ -1,79 +0,0 @@
-"""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()