File size: 4,915 Bytes
5c49e27 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Convert the released MEG-GPT + EphysTokenizer weights into one braindecode MEGGPT state dict.
Inputs (OHBA-analysis/MEG-GPT @78ff9439, models/):
meg-gpt/model_state.pt MEGGPT (meg_gpt/models/meg_gpt.py) state dict
tokenizer/model_state.pt EphysTokenizer state dict
tokenizer/vocab.pkl vocabulary ({"label_map": codebook entry -> token id, ...})
Dropped on purpose (documented in the PR): the Cam-CAN session embedding
(input_embedding_layer.extra_embeds.0.*), the next-token prediction head, the
attention-mask buffers (rebuilt from the config, checked equal in parity.py),
the tokenizer decoder (token_basis_layer, only used to reconstruct signals) and
the annealing temperature buffer.
usage: python convert.py <models dir> <out dir>
<models dir> is ``models/`` of the original repository; <out dir> receives the
braindecode export (``MEGGPT.save_pretrained``: config.json, model.safetensors,
pytorch_model.bin) that ``MEGGPT.from_pretrained(<out dir or Hub repo>)`` loads.
The classification head (n_outputs=2) is not part of the release: it is
randomly initialised (seed 0); pass ``n_outputs`` to ``from_pretrained`` to
replace it.
"""
import pickle
import re
import sys
from pathlib import Path
import torch
_GPT_RULES = [
(r"input_embedding_layer\.token_embed\.base_module\.weight", "token_embedding.weight"),
(r"input_embedding_layer\.pos_embed\.base_module\.position_embeddings", "time_embedding"),
(r"input_embedding_layer\.channel_embed\.base_module\.position_embeddings", "channel_embedding"),
(r"transformer_decoder\.layers\.(\d+)\.norm_time1\.norm_layer\.", r"layers.\1.norm1."),
(r"transformer_decoder\.layers\.(\d+)\.norm_time2\.norm_layer\.", r"layers.\1.norm2."),
(r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patch_projection_k\.", r"layers.\1.attn.patch_k."),
(r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patch_projection_v\.", r"layers.\1.attn.patch_v."),
(r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patched_projection_k\.", r"layers.\1.attn.k_proj."),
(r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.patched_projection_v\.", r"layers.\1.attn.v_proj."),
(r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.query_projection\.", r"layers.\1.attn.q_proj."),
(r"transformer_decoder\.layers\.(\d+)\.gasp_time_attention\.output_projection\.", r"layers.\1.attn.out_proj."),
(r"transformer_decoder\.layers\.(\d+)\.ff_time\.net\.0\.", r"layers.\1.mlp.0."),
(r"transformer_decoder\.layers\.(\d+)\.ff_time\.net\.2\.", r"layers.\1.mlp.3."),
]
_GPT_DROP = re.compile(
r"input_embedding_layer\.extra_embeds\.|prediction_head\.|.*\.gasp_layer\.attention_mask$"
)
_TOK_RULES = [
(r"encoder_layer\.rnn_layers\.0\.", "tokenizer.rnn."),
(r"token_weights_layer\.dense_layer\.", "tokenizer.proj."),
(r"token_weights_layer\.norm_layer\.", "tokenizer.norm."),
]
_TOK_DROP = re.compile(r"decoder_layer\.token_basis_layer\.|token_weights_layer\._temperature$")
def _rename(state, rules, drop):
out, dropped = {}, []
for key, value in state.items():
if drop.match(key):
dropped.append(key)
continue
for pattern, repl in rules:
new, n = re.subn("^" + pattern, repl, key)
if n:
out[new] = value
break
else:
raise KeyError(f"no conversion rule for {key}")
return out, dropped
def convert(gpt_state, tok_state, vocab):
"""Return (braindecode state dict, list of dropped original keys)."""
gpt, dropped_gpt = _rename(gpt_state, _GPT_RULES, _GPT_DROP)
tok, dropped_tok = _rename(tok_state, _TOK_RULES, _TOK_DROP)
state = {**tok, **gpt}
state["tokenizer.vocab"] = torch.as_tensor(vocab["label_map"], dtype=torch.long)
return state, dropped_gpt + dropped_tok
def load_released(models_dir):
models_dir = Path(models_dir)
gpt = torch.load(models_dir / "meg-gpt/model_state.pt", map_location="cpu", weights_only=True)
tok = torch.load(models_dir / "tokenizer/model_state.pt", map_location="cpu", weights_only=True)
with open(models_dir / "tokenizer/vocab.pkl", "rb") as f:
vocab = pickle.load(f)
return gpt, tok, vocab
def export(models_dir, out_dir):
from braindecode.models import MEGGPT
state, dropped = convert(*load_released(models_dir))
torch.manual_seed(0)
model = MEGGPT(n_chans=52, n_outputs=2, sfreq=250).eval()
res = model.load_state_dict(state, strict=False)
assert sorted(res.missing_keys) == ["final_layer.1.bias", "final_layer.1.weight"], res
assert not res.unexpected_keys, res
model.save_pretrained(out_dir)
return state, dropped
if __name__ == "__main__":
state, dropped = export(sys.argv[1], sys.argv[2])
print(f"exported {len(state)} converted tensors to {sys.argv[2]}; dropped {len(dropped)}: {dropped}")
|