Add MEG-XL meg-xl-med-v2 weights converted for braindecode.models.MEGXL (pnpl/MEG-XL@79838987, braindecode PR #1554 @ bf95a14d)
5b5dcbc verified | """Convert the released MEG-XL Lightning checkpoint to braindecode's MEGXL. | |
| usage: python convert.py meg-xl-med-v2.ckpt <out_dir> <fif with the default chs_info> | |
| The checkpoint holds the frozen BioCodec tokenizer, so ``biocodec_ckpt.pt`` is | |
| not needed (its tensors equal the copy inside the checkpoint, max-abs 0.0). | |
| The keys are renamed to braindecode's (``RENAME``). Dropped: the BioCodec | |
| decoder, the duplicated ``biocodec_embedding_*`` codebooks, ``mask_token``, | |
| ``output_head`` (pretraining only) and the complex RoPE cache. The port's | |
| RoPE table (``rotate``, as many positions as the model's buffer: 625, the | |
| 150 s pretraining window) is recomputed from ``freqs`` with ``torch.polar`` | |
| on the CPU, as the authors' loader recomputes it, and stored as ``(cos, sin)`` | |
| pairs. | |
| ``<out_dir>`` gets ``config.json`` from ``save_pretrained`` (the 306 MEG | |
| channels of the given FIF file as the default ``chs_info``, 150 s at 50 Hz) | |
| and ``model.safetensors`` / ``pytorch_model.bin`` without the classification | |
| head: its shape depends on the number of channels and it is not pretrained, | |
| so ``from_pretrained`` builds a new one for the caller's ``chs_info``. | |
| """ | |
| import sys | |
| import torch | |
| DROP = ("tokenizer.decoder.", "biocodec_embedding_", "mask_token", "output_head.") | |
| RENAME = ( | |
| ("tokenizer.encoder.model.", "encoder."), | |
| (".conv.conv.", ".conv."), | |
| ("tokenizer.quantizer.vq.", "quantizer."), | |
| ("criss_cross_transformer.layers.", "blocks."), | |
| (".ff.layer.0.", ".ff.0."), | |
| (".ff.layer.2.", ".ff.3."), | |
| ("position_fourier_emb.B", "position_freqs"), | |
| ("orientation_fourier_emb.B", "orientation_freqs"), | |
| ) | |
| def convert(state_dict, model): | |
| """Map ``state_dict`` onto ``model``'s keys.""" | |
| out = {} | |
| for key, value in state_dict.items(): | |
| if key.startswith(DROP) or key.endswith("rope_embedding_layer.rotate"): | |
| continue | |
| for old, new in RENAME: | |
| key = key.replace(old, new) | |
| out[key] = value | |
| for key, value in model.state_dict().items(): | |
| if key.endswith("rope_embedding_layer.rotate"): | |
| freqs = out[key.replace("rotate", "freqs")] | |
| angles = torch.outer(torch.arange(value.shape[0]).type_as(freqs), freqs) | |
| out[key] = torch.view_as_real(torch.polar(torch.ones_like(angles), angles)) | |
| return out | |
| if __name__ == "__main__": | |
| import mne | |
| from safetensors.torch import save_file | |
| from braindecode.models import MEGXL | |
| ckpt_path, out_dir, fif = sys.argv[1:4] | |
| raw = mne.io.read_raw_fif(fif, verbose=False) | |
| raw.pick(mne.pick_types(raw.info, meg=True, ref_meg=False, exclude=[])) | |
| torch.manual_seed(0) | |
| model = MEGXL(chs_info=raw.info["chs"], n_outputs=2, n_times=7500, sfreq=50.0) | |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) | |
| state = convert(ckpt["state_dict"], model) | |
| missing, unexpected = model.load_state_dict(state, strict=False) | |
| assert unexpected == [] and set(missing) == {"final_layer.weight", "final_layer.bias"} | |
| model.save_pretrained(out_dir) | |
| weights = { | |
| k: v.contiguous() | |
| for k, v in model.state_dict().items() | |
| if not k.startswith("final_layer.") | |
| } | |
| save_file(weights, f"{out_dir}/model.safetensors", metadata={"format": "pt"}) | |
| torch.save(weights, f"{out_dir}/pytorch_model.bin") | |