File size: 3,316 Bytes
5b5dcbc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """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")
|