Add BBL CamCAN weights converted for braindecode.models.BBL (the-brains-bitter-lesson@3965a2be README checkpoint, braindecode PR #1553 @ 388ac6cb)
5a19a36 verified | """Convert the released BBL Lightning checkpoint (CamCAN Pretrainer of | |
| neural-processing-lab/the-brains-bitter-lesson) to braindecode's ``BBL``. | |
| usage: python convert.py <bbl.ckpt> <out_dir> | |
| Keeps the backbone only (the SSL projector and the band/amp/phase pretext heads | |
| are training-only), renames the keys to braindecode's, and fills the | |
| ``trained_subjects`` buffer with the subject rows the optimizer updated | |
| (non-zero Adam second moment), since the original keeps that set in a Python | |
| attribute that is not saved. | |
| ``<out_dir>`` gets ``config.json``, ``model.safetensors`` and | |
| ``pytorch_model.bin`` from ``BBL.save_pretrained`` for | |
| ``BBL(n_chans=306, n_outputs=2, n_times=125, sfreq=250.0)``. The linear head | |
| (``final_layer``) is not part of the release: it is PyTorch's default init | |
| under ``torch.manual_seed(0)``; ``from_pretrained(..., n_outputs=k)`` replaces it. | |
| """ | |
| import sys | |
| import torch | |
| RENAMES = [ | |
| ("backbone.dataset_gating.projector.", "projector."), | |
| ("backbone.encoder.model.", "encoder."), | |
| ("backbone.subject_embedding.", "subject_embedding."), | |
| (".conv.conv.weight_g", ".conv.parametrizations.weight.original0"), | |
| (".conv.conv.weight_v", ".conv.parametrizations.weight.original1"), | |
| (".conv.conv.bias", ".conv.bias"), | |
| ] | |
| def convert_state_dict(state_dict): | |
| out = {} | |
| for key, value in state_dict.items(): | |
| if not key.startswith("backbone."): | |
| continue | |
| for old, new in RENAMES: | |
| key = key.replace(old, new) | |
| out[key] = value | |
| return out | |
| def trained_subject_mask(checkpoint): | |
| # Adam state 0 is the subject embedding, the first parameter of the Pretrainer. | |
| state = checkpoint["optimizer_states"][0]["state"][0]["exp_avg_sq"] | |
| assert state.shape == checkpoint["state_dict"]["backbone.subject_embedding.weight"].shape | |
| return (state.abs().sum(dim=1) > 0).float() | |
| def convert(checkpoint): | |
| out = convert_state_dict(checkpoint["state_dict"]) | |
| out["trained_subjects"] = trained_subject_mask(checkpoint) | |
| return out | |
| if __name__ == "__main__": | |
| from braindecode.models import BBL | |
| ckpt_path, out_dir = sys.argv[1:3] | |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) | |
| state = convert(ckpt) | |
| torch.manual_seed(0) | |
| model = BBL(n_chans=306, n_outputs=2, n_times=125, sfreq=250.0).eval() | |
| result = model.load_state_dict(state, strict=False) | |
| assert sorted(result.missing_keys) == ["final_layer.bias", "final_layer.weight"], result | |
| assert not result.unexpected_keys, result | |
| model.save_pretrained(out_dir) | |
| print("hyper_parameters", ckpt["hyper_parameters"]) | |
| print("n_keys", len(state), "trained subjects", int(state["trained_subjects"].sum())) | |