# type: ignore
""".. _load-pretrained-models:
Loading and Adapting Pretrained Foundation Models
=================================================
.. raw:: html
+
All braindecode models can load and save weights on the
`Hugging Face Hub `_ via
``from_pretrained`` and ``push_to_hub``. For **foundation models**
we additionally provide curated pretrained checkpoints with mapped
weights, so you can fine-tune them out of the box.
.. important::
The curation and standardization of the pretrained foundation model
weights available on the Braindecode Hugging Face organization was
carried out as part of the
`OpenEEG-Bench `_
benchmark [2]_.
If you use these pretrained weights, please cite the paper.
.. code-block:: bash
pip install braindecode[hub]
This tutorial shows how to load pretrained EEG foundation models,
adapt them to new tasks, extract features, and save/restore full model
configurations, using a **unified API** inspired by the
`Hugging Face transformers `_
library.
.. contents:: This example covers:
:local:
:depth: 2
"""
# Authors: Bruno Aristimunha
#
# License: BSD (3-clause)
import json
import os
import warnings
import torch
from huggingface_hub import login
warnings.simplefilter("ignore")
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
login(token=hf_token)
######################################################################
# Loading a pretrained model
# --------------------------
#
# All braindecode foundation models support the ``from_pretrained``
# method, which downloads model weights and configuration from the
# Hugging Face Hub.
#
# We start with BENDR [1]_ as an example:
from braindecode.models import BENDR
model = BENDR.from_pretrained("braindecode/braindecode-bendr", n_outputs=2)
print(f"Loaded BENDR with n_outputs={model.n_outputs}")
######################################################################
# The loaded model is ready for inference:
x = torch.randn(2, 20, 768) # (batch, n_chans, n_times)
model.eval()
with torch.no_grad():
out = model(x)
print(f"Output shape: {out.shape}")
######################################################################
# Adapting to a new task
# ----------------------
#
# Changing the number of outputs
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# When fine-tuning a pretrained model on a dataset with a different number
# of classes, pass ``n_outputs`` directly to ``from_pretrained``. The
# backbone weights are loaded and the classification head is automatically
# rebuilt:
model = BENDR.from_pretrained("braindecode/braindecode-bendr", n_outputs=4)
print(f"n_outputs after loading: {model.n_outputs}")
with torch.no_grad():
out = model(x)
print(f"Output shape with 4 classes: {out.shape}")
######################################################################
# You can also swap the head at any time using ``reset_head``:
model.reset_head(10)
print(f"n_outputs after reset_head: {model.n_outputs}")
with torch.no_grad():
out = model(x)
print(f"Output shape after reset: {out.shape}")
######################################################################
# Extracting features
# -------------------
#
# All foundation models support ``return_features=True`` in their
# ``forward()`` method. This returns a dictionary with:
#
# - ``"features"`` -- encoder embeddings before the classification head
# - ``"cls_token"`` -- the CLS token (if the model has one, otherwise ``None``)
model.eval()
with torch.no_grad():
out = model(x, return_features=True)
print(f"Type: {type(out)}")
print(f"Features shape: {out['features'].shape}")
print(f"CLS token: {out['cls_token']}")
######################################################################
# .. tip::
#
# This is useful for **transfer learning**: freeze the backbone and
# train only a new head on the extracted features.
# See :ref:`finetune-foundation-model` for a complete example.
######################################################################
# Saving and restoring configurations
# ------------------------------------
#
# ``get_config`` returns a JSON-serializable dictionary of **all**
# ``__init__`` parameters (not just the 6 EEG-specific ones). This
# includes model-specific hyperparameters like ``encoder_h``,
# ``drop_prob``, ``activation``, etc.
config = model.get_config()
print(json.dumps({k: v for k, v in config.items() if k != "chs_info"}, indent=2))
######################################################################
# You can reconstruct the model (without weights) using ``from_config``:
model_copy = BENDR.from_config(config)
print(f"Reconstructed: n_outputs={model_copy.n_outputs}")
######################################################################
# When pushing to the Hub, the full config is saved automatically in
# ``config.json`` alongside the model weights:
#
# .. code-block:: python
#
# model.push_to_hub("username/my-bendr-model")
# # Later:
# model = BENDR.from_pretrained("username/my-bendr-model")
######################################################################
# Unified API across foundation models
# -------------------------------------
#
# The same API works across **all** foundation models:
#
# .. list-table::
# :header-rows: 1
# :widths: 25 15 15 20 15
#
# * - Model
# - ``from_pretrained``
# - ``reset_head``
# - ``return_features``
# - ``get_config``
# * - :class:`~braindecode.models.BENDR`
# - |check|
# - |check|
# - |check|
# - |check|
# * - :class:`~braindecode.models.BIOT`
# - |check|
# - |check|
# - |check|
# - |check|
# * - :class:`~braindecode.models.CBraMod`
# - |check|
# - |check|
# - |check|
# - |check|
# * - :class:`~braindecode.models.EEGPT`
# - |check|
# - |check|
# - |check|
# - |check|
# * - :class:`~braindecode.models.Labram`
# - |check|
# - |check|
# - |check| (+ ``cls_token``)
# - |check|
# * - :class:`~braindecode.models.LUNA`
# - |check|
# - |check|
# - |check|
# - |check|
# * - :class:`~braindecode.models.REVE`
# - |check|
# - |check|
# - |check|
# - |check|
# * - :class:`~braindecode.models.SignalJEPA` variants
# - |check|
# - |check|
# - |check|
# - |check|
#
# .. |check| unicode:: 0x2714
#
# The feature shapes differ between models (reflecting their
# architecture), but the API is always the same.
######################################################################
# Available pretrained weights
# ----------------------------
#
# .. raw:: html
#
#
#
#