You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
"\n# Fine-tuning a Foundation Model (Signal-JEPA)\n\nFoundation models are large-scale pre-trained models that serve as a starting point\nfor a wide range of downstream tasks, leveraging their generalization capabilities.\nFine-tuning these models is necessary to adapt them to specific tasks or datasets,\nensuring optimal performance in specialized applications.\n\nIn this tutorial, we demonstrate how to load a pre-trained foundation model\nand fine-tune it for a specific task. We use the Signal-JEPA model [1]_\nand a MOABB motor-imagery dataset for this tutorial.\n :depth: 2\n"
"## Loading and preparing the data\n\n### Loading a dataset\n\nWe start by loading a MOABB dataset, a single subject only for speed.\nThe dataset contains motor imagery EEG recordings, which we will preprocess and use for fine-tuning.\n\n\n"
26
+
]
27
+
},
28
+
{
29
+
"cell_type": "code",
30
+
"execution_count": null,
31
+
"metadata": {
32
+
"collapsed": false
33
+
},
34
+
"outputs": [],
35
+
"source": [
36
+
"subject_id = 3 # Just one subject for speed\ndataset = MOABBDataset(dataset_name=\"BNCI2014_001\", subject_ids=[subject_id])\n\n# Set the standard 10-20 montage for EEG channel locations\nmontage = mne.channels.make_standard_montage(\"standard_1020\")\nfor ds in dataset.datasets:\n ds.raw.set_montage(montage)"
37
+
]
38
+
},
39
+
{
40
+
"cell_type": "markdown",
41
+
"metadata": {},
42
+
"source": [
43
+
"### Define Dataset parameters\n\nWe extract the sampling frequency and ensure that it is consistent across\nall recordings. We also extract the window size from the annotations and\ninformation about the EEG channels (names, positions, etc.).\n\n\n"
44
+
]
45
+
},
46
+
{
47
+
"cell_type": "code",
48
+
"execution_count": null,
49
+
"metadata": {
50
+
"collapsed": false
51
+
},
52
+
"outputs": [],
53
+
"source": [
54
+
"# Extract sampling frequency\nsfreq = dataset.datasets[0].raw.info[\"sfreq\"]\nassert all([ds.raw.info[\"sfreq\"] == sfreq for ds in dataset.datasets])\n\n# Extract and validate window size from annotations\nwindow_size_seconds = dataset.datasets[0].raw.annotations.duration[0]\nassert all(\n d == window_size_seconds\n for ds in dataset.datasets\n for d in ds.raw.annotations.duration\n)\n\n# Extract channel information\nchs_info = dataset.datasets[0].raw.info[\"chs\"] # Channel information\n\nprint(f\"{sfreq=}, {window_size_seconds=}, {len(chs_info)=}\")"
55
+
]
56
+
},
57
+
{
58
+
"cell_type": "markdown",
59
+
"metadata": {},
60
+
"source": [
61
+
"### Create Windows from Events\n\nWe use the `create_windows_from_events` function from Braindecode to segment\nthe dataset into windows based on events.\n\n\n"
62
+
]
63
+
},
64
+
{
65
+
"cell_type": "code",
66
+
"execution_count": null,
67
+
"metadata": {
68
+
"collapsed": false
69
+
},
70
+
"outputs": [],
71
+
"source": [
72
+
"classes = [\"feet\", \"left_hand\", \"right_hand\"]\nclasses_mapping = {c: i for i, c in enumerate(classes)}\n\nwindows_dataset = create_windows_from_events(\n dataset,\n preload=True, # Preload the data into memory for faster processing\n mapping=classes_mapping,\n)\nmetadata = windows_dataset.get_metadata()\nprint(metadata.head(10))"
73
+
]
74
+
},
75
+
{
76
+
"cell_type": "markdown",
77
+
"metadata": {},
78
+
"source": [
79
+
"## Loading a pre-trained foundation model\n\n### Download and Load Pre-trained Weights\n\nWe download the pre-trained weights for the SignalJEPA model from the Hugging Face Hub.\nThese weights will serve as the starting point for finetuning.\n\n\n"
"### Instantiate the Foundation Model\n\nWe create an instance of the SignalJEPA model using the pre-local downstream\narchitecture. The model is initialized with the dataset's sampling frequency,\nwindow size, and channel information.\n\n\n"
"### Load the Pre-trained Weights into the Model\n\nWe load the pre-trained weights into the model. The transformer layers are excluded\nas this module is not used in the pre-local downstream architecture (see [1]_).\n\n\n"
116
+
]
117
+
},
118
+
{
119
+
"cell_type": "code",
120
+
"execution_count": null,
121
+
"metadata": {
122
+
"collapsed": false
123
+
},
124
+
"outputs": [],
125
+
"source": [
126
+
"# Define layers to exclude from the pre-trained weights\nnew_layers = {\n\"spatial_conv.1.weight\",\n\"spatial_conv.1.bias\",\n\"final_layer.1.weight\",\n\"final_layer.1.bias\",\n}\n\n# Filter out transformer weights and load the state dictionary\nmodel_state_dict = {\n k: v for k, v in model_state_dict.items() if not k.startswith(\"transformer.\")\n}\nmissing_keys, unexpected_keys = model.load_state_dict(model_state_dict, strict=False)\n\n# Ensure no unexpected keys and validate missing keys\nassert unexpected_keys == [], f\"{unexpected_keys=}\"\nassert set(missing_keys) == new_layers, f\"{missing_keys=}\""
127
+
]
128
+
},
129
+
{
130
+
"cell_type": "markdown",
131
+
"metadata": {},
132
+
"source": [
133
+
"## Fine-tuning the Model\n\nSignal-JEPA is a model trained in a self-supervised manner on a masked\nprediction task. In this task, the model is configured in a many-to-many\nfashion, which is not suited for a classification task. Therefore, we need to\nadjust the model architecture for finetuning. This is what is done by the\n:class:`SignalJEPA_PreLocal`, :class:`SignalJEPA_Contextual`, and\n:class:`SignalJEPA_PostLocal` classes. In these classes, new layers are added\nspecifically for classification, as described in the article [1]_ and in the following figure:\n\n<img src=\"file://_static/model/sjepa_pre-local.jpg\" alt=\"Signal-JEPA Pre-Local Downstream Architecture\" align=\"center\">\n\nWith this downstream architecture, two options are possible for fine-tuning:\n\n1) Fine-tune only the newly added layers\n2) Fine-tune the entire model\n\n### Freezing Pre-trained Layers\n\nAs the second option is rather straightforward to implement,\nwe will focus on the first option here.\nWe will freeze all layers except the newly added ones.\n\n\n"
134
+
]
135
+
},
136
+
{
137
+
"cell_type": "code",
138
+
"execution_count": null,
139
+
"metadata": {
140
+
"collapsed": false
141
+
},
142
+
"outputs": [],
143
+
"source": [
144
+
"for name, param in model.named_parameters():\n if name not in new_layers:\n param.requires_grad = False\n\nprint(\"Trainable parameters:\")\nother_modules = set()\nfor name, param in model.named_parameters():\n if param.requires_grad:\n print(name)\n else:\n other_modules.add(name.split(\".\")[0])\n\nprint(\"\\nOther modules:\")\nprint(other_modules)"
145
+
]
146
+
},
147
+
{
148
+
"cell_type": "markdown",
149
+
"metadata": {},
150
+
"source": [
151
+
"### Fine-tuning Procedure\n\nFinally, we set up the fine-tuning procedure using Braindecode's\n:class:`EEGClassifier`. We define the loss function, optimizer, and training\nparameters. We then fit the model to the windows dataset.\n\nWe only train for a few epochs for demonstration purposes.\n\n\n"
"### All-in-one Implementation\n\nIn the implementation above, we manually loaded the weights and froze the layers.\nThis forces us to pass an initialized model to :class:`EEGClassifier`, which may\ncreate issues if we use it in a cross-validation setting.\n\nInstead, we can implement the same procedure in a more compact and reproducible way,\nby using skorch's callback system.\n\nHere, we import a callback to freeze layers and define a custom\ncallback to load the pre-trained weights at the beginning of training:\n\n\n"
"## Conclusion and Next Steps\n\nIn this tutorial, we demonstrated how to fine-tune a pre-trained foundation\nmodel, Signal-JEPA, for a motor imagery classification task. We now have a basic\nimplementation that can automatically load pre-trained weights and freeze specific layers.\n\nThis setup can easily be extended to explore different fine-tuning techniques,\nbase foundation models, and downstream tasks.\n\n\n"
206
+
]
207
+
},
208
+
{
209
+
"cell_type": "markdown",
210
+
"metadata": {},
211
+
"source": [
212
+
"## References\n\n.. [1] Guetschel, P., Moreau, T., and Tangermann, M. (2024)\n\u201cS-JEPA: towards seamless cross-dataset transfer\n through dynamic spatial attention\u201d. https://arxiv.org/abs/2403.11772\n\n"
0 commit comments