File size: 4,438 Bytes
d111525 | 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 108 109 110 | """Convert the released BrainOmni checkpoints to braindecode-native files.
Source: https://huggingface.co/OpenTSLab/BrainOmni at revision
9a4d3c70495370397ccfbfd6d2496f25647545a5 (MIT), files
``braintokenizer/BrainTokenizer.pt``, ``tiny/BrainOmni.pt``, ``base/BrainOmni.pt``
and their ``model_cfg.json``.
Usage::
python convert_brainomni_checkpoints.py OUT_DIR [SOURCE_DIR]
writes ``OUT_DIR/{braintokenizer,brainomni-tiny,brainomni-base}-pretrained``
with ``config.json``, ``model.safetensors`` and ``pytorch_model.bin``
(``save_pretrained``). Without ``SOURCE_DIR`` the files are downloaded.
Key changes: ``quantizer.rvq.`` -> ``quantizer.``, ``decoder.`` ->
``final_layer.``, the doubled ``conv``/``convtr`` wrappers and SEANet's
``.model`` are flattened, ``weight_g``/``weight_v`` -> weight-norm
parametrizations, the feed-forward linears are renumbered as
``FeedForwardBlock``. The pretraining-only ``mask_token`` and ``predict_head``
are dropped. The RoPE cache ``rotate`` is stored by the release as cosines only
(the released code copies it into a complex cache, so the sine is zero); it is
kept as ``(cos, 0)`` pairs. The BrainOmni classification head is not
pretrained: it is a seeded random ``nn.Linear`` default init.
The stored ``chs_info`` (19 EEG channels, 10-20) is only a default; pass
your own ``chs_info`` to ``from_pretrained``.
"""
import json
import sys
from pathlib import Path
import mne
import torch
from braindecode.models import BrainOmni, BrainTokenizer
REPO, REVISION = "OpenTSLab/BrainOmni", "9a4d3c70495370397ccfbfd6d2496f25647545a5"
RENAMES = {
"n_dim": "emb_dim",
"n_head": "tokenizer_num_heads",
"lm_head": "num_heads",
"lm_depth": "depth",
"lm_dropout": "drop_prob",
}
UNUSED = {"mask_ratio", "num_quantizers_used", "quantize_optimize_method"}
CH_NAMES = (
"Fp1 Fp2 F7 F3 Fz F4 F8 T7 C3 Cz C4 T8 P7 P3 Pz P4 P8 O1 O2".split()
)
def fetch(source, path):
if source is not None:
return Path(source) / path
from huggingface_hub import hf_hub_download
return Path(hf_hub_download(REPO, path, revision=REVISION))
def rename(key):
key = key.replace("quantizer.rvq.", "quantizer.")
if key.startswith("decoder."):
key = "final_layer." + key.removeprefix("decoder.")
key = key.replace("tokenizer.decoder.", "tokenizer.final_layer.")
key = key.replace(".convtr.convtr.", ".convtr.").replace(".conv.conv.", ".conv.")
key = key.replace("seanet_encoder.model.", "seanet_encoder.")
key = key.replace("seanet_decoder.model.", "seanet_decoder.")
key = key.replace(".weight_g", ".parametrizations.weight.original0")
key = key.replace(".weight_v", ".parametrizations.weight.original1")
for name in ("ff", "aggregate_mlp"):
key = key.replace(f"{name}.layer.0.", f"{name}.0.")
key = key.replace(f"{name}.layer.2.", f"{name}.3.")
return key
def convert(source, folder, cls, out, **kwargs):
config = json.loads(fetch(source, f"{folder}/model_cfg.json").read_text())
config = {RENAMES.get(k, k): v for k, v in config.items() if k not in UNUSED}
dropout = config.pop("dropout")
config["drop_prob" if cls is BrainTokenizer else "tokenizer_drop_prob"] = dropout
info = mne.create_info(CH_NAMES, 256.0, "eeg")
info.set_montage("standard_1020")
torch.manual_seed(0) # the BrainOmni head is a seeded random init
model = cls(chs_info=info["chs"], n_times=512, sfreq=256.0, **config, **kwargs)
official = torch.load(
fetch(source, f"{folder}/{cls.__name__}.pt"), weights_only=True
)
state = {}
for key, value in official.items():
if key == "mask_token" or key.startswith("predict_head."):
continue
if key.endswith("rope_embedding_layer.rotate"):
value = torch.stack((value, torch.zeros_like(value)), dim=-1)
state[rename(key)] = value
for key, value in model.state_dict().items():
if key.startswith("final_layer.") and cls is BrainOmni:
state[key] = value
model.load_state_dict(state, strict=True)
model.save_pretrained(out)
return model
if __name__ == "__main__":
out = Path(sys.argv[1])
source = sys.argv[2] if len(sys.argv) > 2 else None
convert(source, "braintokenizer", BrainTokenizer, out / "braintokenizer-pretrained")
for size in ("tiny", "base"):
convert(source, size, BrainOmni, out / f"brainomni-{size}-pretrained", n_outputs=2)
|