File size: 2,715 Bytes
5a19a36 | 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 | """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()))
|