braindecode 1.6.0.dev1028__py3-none-any.whl → 1.6.1.dev1014__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- braindecode/datasets/bids/hub.py +46 -1
- braindecode/datasets/moabb.py +0 -10
- braindecode/eegneuralnet.py +41 -0
- braindecode/models/__init__.py +7 -3
- braindecode/models/atcnet.py +118 -1
- braindecode/models/attn_sleep.py +0 -10
- braindecode/models/eegdino.py +533 -0
- braindecode/models/eegnet.py +1 -10
- braindecode/models/eegpt.py +29 -0
- braindecode/models/eegtcnet.py +62 -27
- braindecode/models/sstdpn.py +10 -4
- braindecode/models/summary.csv +55 -52
- braindecode/models/tcformer.py +643 -0
- braindecode/models/tsinception.py +1 -11
- braindecode/models/util.py +7 -2
- braindecode/modules/__init__.py +2 -1
- braindecode/modules/blocks.py +82 -0
- braindecode/preprocessing/mne_preprocess.py +28 -1
- braindecode/version.py +1 -1
- {braindecode-1.6.0.dev1028.dist-info → braindecode-1.6.1.dev1014.dist-info}/METADATA +2 -2
- {braindecode-1.6.0.dev1028.dist-info → braindecode-1.6.1.dev1014.dist-info}/RECORD +25 -23
- {braindecode-1.6.0.dev1028.dist-info → braindecode-1.6.1.dev1014.dist-info}/WHEEL +0 -0
- {braindecode-1.6.0.dev1028.dist-info → braindecode-1.6.1.dev1014.dist-info}/licenses/LICENSE.txt +0 -0
- {braindecode-1.6.0.dev1028.dist-info → braindecode-1.6.1.dev1014.dist-info}/licenses/NOTICE.txt +0 -0
- {braindecode-1.6.0.dev1028.dist-info → braindecode-1.6.1.dev1014.dist-info}/top_level.txt +0 -0
braindecode/datasets/bids/hub.py
CHANGED
|
@@ -327,6 +327,7 @@ class HubDatasetMixin:
|
|
|
327
327
|
|
|
328
328
|
n_recordings = len(self.datasets)
|
|
329
329
|
first_ds = self.datasets[0]
|
|
330
|
+
descriptions = [ds.description for ds in self.datasets]
|
|
330
331
|
|
|
331
332
|
# Get dataset-specific info based on type using registry
|
|
332
333
|
dataset_type = get_dataset_type(first_ds)
|
|
@@ -358,6 +359,9 @@ class HubDatasetMixin:
|
|
|
358
359
|
raise TypeError(f"Unsupported dataset type: {dataset_type}")
|
|
359
360
|
|
|
360
361
|
# Create README content and save
|
|
362
|
+
license_value, license_needs_update = _infer_license_from_descriptions(
|
|
363
|
+
descriptions
|
|
364
|
+
)
|
|
361
365
|
readme_content = _generate_readme_content(
|
|
362
366
|
format_info=format_info,
|
|
363
367
|
n_recordings=n_recordings,
|
|
@@ -366,6 +370,8 @@ class HubDatasetMixin:
|
|
|
366
370
|
data_type=data_type,
|
|
367
371
|
n_windows=n_windows,
|
|
368
372
|
total_duration=total_duration,
|
|
373
|
+
license=license_value,
|
|
374
|
+
license_needs_update=license_needs_update,
|
|
369
375
|
)
|
|
370
376
|
|
|
371
377
|
# Save README
|
|
@@ -466,6 +472,7 @@ class HubDatasetMixin:
|
|
|
466
472
|
token: Optional[str] = None,
|
|
467
473
|
cache_dir: Optional[Union[str, Path]] = None,
|
|
468
474
|
force_download: bool = False,
|
|
475
|
+
revision: str | None = None,
|
|
469
476
|
**kwargs,
|
|
470
477
|
):
|
|
471
478
|
"""
|
|
@@ -485,6 +492,9 @@ class HubDatasetMixin:
|
|
|
485
492
|
cache directory (~/.cache/huggingface/datasets).
|
|
486
493
|
force_download : bool, default=False
|
|
487
494
|
Whether to force re-download even if cached.
|
|
495
|
+
revision : str | None, default=None
|
|
496
|
+
Specific branch, tag, or commit to download. If None, uses the
|
|
497
|
+
repository's default revision.
|
|
488
498
|
**kwargs
|
|
489
499
|
Additional arguments (currently unused).
|
|
490
500
|
|
|
@@ -526,6 +536,7 @@ class HubDatasetMixin:
|
|
|
526
536
|
token=token,
|
|
527
537
|
cache_dir=cache_dir,
|
|
528
538
|
force_download=force_download,
|
|
539
|
+
revision=revision,
|
|
529
540
|
)
|
|
530
541
|
|
|
531
542
|
# Load format info
|
|
@@ -962,6 +973,8 @@ def _generate_readme_content(
|
|
|
962
973
|
n_windows: int,
|
|
963
974
|
total_duration: float | None = None,
|
|
964
975
|
format: str = "zarr",
|
|
976
|
+
license: str = "please-specify",
|
|
977
|
+
license_needs_update: bool = True,
|
|
965
978
|
):
|
|
966
979
|
"""Generate README.md content for a dataset uploaded to the Hub.
|
|
967
980
|
|
|
@@ -983,6 +996,10 @@ def _generate_readme_content(
|
|
|
983
996
|
Total duration in seconds across all recordings.
|
|
984
997
|
format : str
|
|
985
998
|
Storage format (default: "zarr").
|
|
999
|
+
license : str
|
|
1000
|
+
Dataset license string to write in YAML front matter.
|
|
1001
|
+
license_needs_update : bool
|
|
1002
|
+
Whether to append a reminder comment next to the license field.
|
|
986
1003
|
|
|
987
1004
|
Returns
|
|
988
1005
|
-------
|
|
@@ -998,6 +1015,10 @@ def _generate_readme_content(
|
|
|
998
1015
|
str(timedelta(seconds=int(total_duration))) if total_duration else "N/A"
|
|
999
1016
|
)
|
|
1000
1017
|
|
|
1018
|
+
license_line = f"license: {license}"
|
|
1019
|
+
if license_needs_update:
|
|
1020
|
+
license_line += " # Please update this field to reflect your dataset's license"
|
|
1021
|
+
|
|
1001
1022
|
return f"""---
|
|
1002
1023
|
tags:
|
|
1003
1024
|
- braindecode
|
|
@@ -1005,7 +1026,7 @@ tags:
|
|
|
1005
1026
|
- neuroscience
|
|
1006
1027
|
- brain-computer-interface
|
|
1007
1028
|
- deep-learning
|
|
1008
|
-
|
|
1029
|
+
{license_line}
|
|
1009
1030
|
---
|
|
1010
1031
|
|
|
1011
1032
|
# EEG Dataset
|
|
@@ -1101,3 +1122,27 @@ if hasattr(dataset.datasets[0], "bids_channels"):
|
|
|
1101
1122
|
|
|
1102
1123
|
*Created with [braindecode](https://braindecode.org)*
|
|
1103
1124
|
"""
|
|
1125
|
+
|
|
1126
|
+
|
|
1127
|
+
def _infer_license_from_descriptions(
|
|
1128
|
+
descriptions: list[pd.Series | None],
|
|
1129
|
+
) -> tuple[str, bool]:
|
|
1130
|
+
"""Infer a single dataset license from descriptions.
|
|
1131
|
+
|
|
1132
|
+
Returns inferred license and whether a manual update reminder is needed.
|
|
1133
|
+
"""
|
|
1134
|
+
licenses = set()
|
|
1135
|
+
for description in descriptions:
|
|
1136
|
+
if description is None:
|
|
1137
|
+
continue
|
|
1138
|
+
raw_license = description.get("license", None)
|
|
1139
|
+
if raw_license is None:
|
|
1140
|
+
continue
|
|
1141
|
+
normalized = str(raw_license).strip()
|
|
1142
|
+
if normalized and normalized.lower() != "nan":
|
|
1143
|
+
licenses.add(normalized)
|
|
1144
|
+
|
|
1145
|
+
if len(licenses) == 1:
|
|
1146
|
+
return next(iter(licenses)), False
|
|
1147
|
+
|
|
1148
|
+
return "please-specify", True
|
braindecode/datasets/moabb.py
CHANGED
|
@@ -15,7 +15,6 @@ from typing import Any
|
|
|
15
15
|
|
|
16
16
|
import mne
|
|
17
17
|
import pandas as pd
|
|
18
|
-
from mne.utils import deprecated
|
|
19
18
|
|
|
20
19
|
from braindecode.util import _update_moabb_docstring
|
|
21
20
|
|
|
@@ -208,12 +207,3 @@ class HGD(MOABBDataset):
|
|
|
208
207
|
|
|
209
208
|
def __init__(self, subject_ids):
|
|
210
209
|
super().__init__("Schirrmeister2017", subject_ids=subject_ids)
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
@deprecated(
|
|
214
|
-
"`BNCI2014001` was renamed to `BNCI2014_001` in v1.13; this alias will be removed in v1.14."
|
|
215
|
-
)
|
|
216
|
-
class BNCI2014001(BNCI2014_001):
|
|
217
|
-
"""Deprecated alias for BNCI2014001."""
|
|
218
|
-
|
|
219
|
-
pass
|
braindecode/eegneuralnet.py
CHANGED
|
@@ -7,6 +7,7 @@
|
|
|
7
7
|
import abc
|
|
8
8
|
import inspect
|
|
9
9
|
import logging
|
|
10
|
+
import warnings
|
|
10
11
|
from typing import Literal
|
|
11
12
|
|
|
12
13
|
import mne
|
|
@@ -41,6 +42,7 @@ def _get_model(model: str):
|
|
|
41
42
|
|
|
42
43
|
class _EEGNeuralNet(NeuralNet, abc.ABC):
|
|
43
44
|
signal_args_set_ = False
|
|
45
|
+
_warned_drop_last_ = False
|
|
44
46
|
|
|
45
47
|
@property
|
|
46
48
|
def log(self):
|
|
@@ -267,6 +269,44 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
|
|
|
267
269
|
X = X.get_data(units="uV")
|
|
268
270
|
return super().get_dataset(X, y)
|
|
269
271
|
|
|
272
|
+
def get_iterator(self, dataset, training=False):
|
|
273
|
+
iterator = super().get_iterator(dataset, training=training)
|
|
274
|
+
# Warn (once per fit) if drop_last would discard every training batch,
|
|
275
|
+
# e.g. a training set smaller than batch_size. This silently skips
|
|
276
|
+
# training without raising any error, which is hard to diagnose.
|
|
277
|
+
if training and not self._warned_drop_last_:
|
|
278
|
+
self._warn_if_training_batches_dropped(iterator)
|
|
279
|
+
return iterator
|
|
280
|
+
|
|
281
|
+
def _warn_if_training_batches_dropped(self, loader):
|
|
282
|
+
"""Warn if ``drop_last`` discards all training batches.
|
|
283
|
+
|
|
284
|
+
With ``iterator_train__drop_last=True`` (the braindecode default), a
|
|
285
|
+
training set smaller than ``batch_size`` yields zero batches, so no
|
|
286
|
+
forward/backward pass runs and the model is left untrained without any
|
|
287
|
+
error. We inspect the resolved training ``DataLoader`` (after
|
|
288
|
+
``train_split`` has been applied), so the count reflects the actual
|
|
289
|
+
number of training examples.
|
|
290
|
+
"""
|
|
291
|
+
self._warned_drop_last_ = True
|
|
292
|
+
if not getattr(loader, "drop_last", False):
|
|
293
|
+
return
|
|
294
|
+
batch_size = getattr(loader, "batch_size", None)
|
|
295
|
+
try:
|
|
296
|
+
n_train = len(loader.dataset)
|
|
297
|
+
except TypeError:
|
|
298
|
+
return # e.g. an IterableDataset whose length is unknown
|
|
299
|
+
if not batch_size or n_train == 0 or n_train >= batch_size:
|
|
300
|
+
return
|
|
301
|
+
warnings.warn(
|
|
302
|
+
f"The training set has {n_train} example(s), which is smaller than "
|
|
303
|
+
f"batch_size={batch_size}. With iterator_train__drop_last=True "
|
|
304
|
+
f"(the braindecode default), all training batches are dropped and "
|
|
305
|
+
f"the model is not trained. Reduce batch_size to at most {n_train}, "
|
|
306
|
+
f"or pass iterator_train__drop_last=False.",
|
|
307
|
+
UserWarning,
|
|
308
|
+
)
|
|
309
|
+
|
|
270
310
|
def partial_fit(self, X, y=None, classes=None, **fit_params):
|
|
271
311
|
"""Fit the module.
|
|
272
312
|
|
|
@@ -366,6 +406,7 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
|
|
|
366
406
|
the module and to the ``self.train_split`` call.
|
|
367
407
|
"""
|
|
368
408
|
# this needs to be executed before the net is initialized:
|
|
409
|
+
self._warned_drop_last_ = False # re-check on every fresh fit
|
|
369
410
|
if not self.signal_args_set_:
|
|
370
411
|
self._set_signal_args(X, y, classes=None)
|
|
371
412
|
self.signal_args_set_ = True
|
braindecode/models/__init__.py
CHANGED
|
@@ -15,13 +15,14 @@ from .deep4 import Deep4Net
|
|
|
15
15
|
from .deepsleepnet import DeepSleepNet
|
|
16
16
|
from .dgcnn import DGCNN
|
|
17
17
|
from .eegconformer import EEGConformer
|
|
18
|
+
from .eegdino import EEGDINO
|
|
18
19
|
from .eeginception_erp import EEGInceptionERP
|
|
19
20
|
from .eeginception_mi import EEGInceptionMI
|
|
20
21
|
from .eegitnet import EEGITNet
|
|
21
22
|
from .eegminer import EEGMiner
|
|
22
|
-
from .eegnet import EEGNet
|
|
23
|
+
from .eegnet import EEGNet
|
|
23
24
|
from .eegnex import EEGNeX
|
|
24
|
-
from .eegpt import EEGPT
|
|
25
|
+
from .eegpt import EEGPT, InterpolatedEEGPT
|
|
25
26
|
from .eegsimpleconv import EEGSimpleConv
|
|
26
27
|
from .eegsym import EEGSym
|
|
27
28
|
from .eegtcnet import EEGTCNet
|
|
@@ -54,6 +55,7 @@ from .sleep_stager_chambon_2018 import SleepStagerChambon2018
|
|
|
54
55
|
from .sparcnet import SPARCNet
|
|
55
56
|
from .sstdpn import SSTDPN
|
|
56
57
|
from .syncnet import SyncNet
|
|
58
|
+
from .tcformer import TCFormer
|
|
57
59
|
from .tcn import BDTCN, TCN
|
|
58
60
|
from .tidnet import TIDNet
|
|
59
61
|
from .tsinception import TSception
|
|
@@ -84,13 +86,14 @@ __all__ = [
|
|
|
84
86
|
"DeepSleepNet",
|
|
85
87
|
"BrainModule",
|
|
86
88
|
"EEGConformer",
|
|
89
|
+
"EEGDINO",
|
|
87
90
|
"EEGPT",
|
|
91
|
+
"InterpolatedEEGPT",
|
|
88
92
|
"EEGInceptionERP",
|
|
89
93
|
"EEGInceptionMI",
|
|
90
94
|
"EEGITNet",
|
|
91
95
|
"EEGMiner",
|
|
92
96
|
"EEGNet",
|
|
93
|
-
"EEGNetv4",
|
|
94
97
|
"EEGNeX",
|
|
95
98
|
"EEGSym",
|
|
96
99
|
"EEGSimpleConv",
|
|
@@ -128,6 +131,7 @@ __all__ = [
|
|
|
128
131
|
"SPARCNet",
|
|
129
132
|
"SyncNet",
|
|
130
133
|
"BDTCN",
|
|
134
|
+
"TCFormer",
|
|
131
135
|
"TCN",
|
|
132
136
|
"TIDNet",
|
|
133
137
|
"TSception",
|
braindecode/models/atcnet.py
CHANGED
|
@@ -2,14 +2,21 @@
|
|
|
2
2
|
#
|
|
3
3
|
# License: BSD (3-clause)
|
|
4
4
|
import math
|
|
5
|
+
from typing import Optional
|
|
5
6
|
|
|
6
7
|
import torch
|
|
7
8
|
from einops.layers.torch import Rearrange
|
|
8
9
|
from mne.utils import warn
|
|
9
10
|
from torch import nn
|
|
11
|
+
from torch.nn.utils.parametrize import is_parametrized, register_parametrization
|
|
10
12
|
|
|
11
13
|
from braindecode.models.base import EEGModuleMixin
|
|
12
|
-
from braindecode.modules import
|
|
14
|
+
from braindecode.modules import (
|
|
15
|
+
CausalConv1d,
|
|
16
|
+
Ensure4d,
|
|
17
|
+
MaxNormLinear,
|
|
18
|
+
MaxNormParametrize,
|
|
19
|
+
)
|
|
13
20
|
|
|
14
21
|
|
|
15
22
|
class ATCNet(EEGModuleMixin, nn.Module):
|
|
@@ -213,6 +220,16 @@ class ATCNet(EEGModuleMixin, nn.Module):
|
|
|
213
220
|
max_norm_const : float
|
|
214
221
|
Maximum L2-norm constraint imposed on weights of the last
|
|
215
222
|
fully-connected layer. Defaults to 0.25.
|
|
223
|
+
conv_max_norm_const : float | None
|
|
224
|
+
If not ``None``, applies a max-norm constraint (via a weight
|
|
225
|
+
parametrization) to the convolution kernels of :class:`_ConvBlock` and
|
|
226
|
+
:class:`_TCNResidualBlock`, matching the official implementation which
|
|
227
|
+
uses ``0.6``. When ``None`` (default), no constraint is applied to those
|
|
228
|
+
layers, preserving the previous behavior. The complementary ``L2``
|
|
229
|
+
weight decay of the official code can be obtained from
|
|
230
|
+
:meth:`source_optimizer_param_groups`.
|
|
231
|
+
|
|
232
|
+
.. versionadded:: 1.6.1
|
|
216
233
|
|
|
217
234
|
Notes
|
|
218
235
|
-----
|
|
@@ -257,6 +274,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
|
|
|
257
274
|
tcn_activation: type[nn.Module] = nn.ELU,
|
|
258
275
|
concat=False,
|
|
259
276
|
max_norm_const=0.25,
|
|
277
|
+
conv_max_norm_const: Optional[float] = None,
|
|
260
278
|
chs_info=None,
|
|
261
279
|
n_times=None,
|
|
262
280
|
):
|
|
@@ -325,6 +343,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
|
|
|
325
343
|
self.tcn_activation = tcn_activation
|
|
326
344
|
self.concat = concat
|
|
327
345
|
self.max_norm_const = max_norm_const
|
|
346
|
+
self.conv_max_norm_const = conv_max_norm_const
|
|
328
347
|
self.tcn_n_filters = int(self.conv_block_depth_mult * self.conv_block_n_filters)
|
|
329
348
|
map = dict()
|
|
330
349
|
for w in range(self.n_windows):
|
|
@@ -346,6 +365,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
|
|
|
346
365
|
pool_size_2=conv_block_pool_size_2,
|
|
347
366
|
depth_mult=conv_block_depth_mult,
|
|
348
367
|
dropout=conv_block_dropout,
|
|
368
|
+
conv_max_norm=self.conv_max_norm_const,
|
|
349
369
|
)
|
|
350
370
|
|
|
351
371
|
self.F2 = int(conv_block_depth_mult * conv_block_n_filters)
|
|
@@ -375,6 +395,7 @@ class ATCNet(EEGModuleMixin, nn.Module):
|
|
|
375
395
|
dropout=self.tcn_dropout,
|
|
376
396
|
activation=self.tcn_activation,
|
|
377
397
|
dilation=2**i,
|
|
398
|
+
conv_max_norm=self.conv_max_norm_const,
|
|
378
399
|
)
|
|
379
400
|
for i in range(self.tcn_depth)
|
|
380
401
|
]
|
|
@@ -458,6 +479,93 @@ class ATCNet(EEGModuleMixin, nn.Module):
|
|
|
458
479
|
|
|
459
480
|
return self.out_fun(sw_concat_agg)
|
|
460
481
|
|
|
482
|
+
def source_optimizer_param_groups(
|
|
483
|
+
self,
|
|
484
|
+
conv_weight_decay: float = 0.009,
|
|
485
|
+
dense_weight_decay: float = 0.5,
|
|
486
|
+
) -> list[dict]:
|
|
487
|
+
r"""Source-faithful optimizer parameter groups for ATCNet.
|
|
488
|
+
|
|
489
|
+
The official Keras implementation applies ``L2`` weight decay only to the
|
|
490
|
+
convolution / TCN kernels (``conv_weightDecay = 0.009``) and to the final
|
|
491
|
+
dense layer (``dense_weightDecay = 0.5``); biases, BatchNorm and attention
|
|
492
|
+
weights are left undecayed. This reproduces the Keras ``L2`` kernel penalty
|
|
493
|
+
only for optimizers that implement ``weight_decay`` as a coupled L2 term,
|
|
494
|
+
e.g. :class:`torch.optim.SGD` or :class:`torch.optim.Adam` (as used by the
|
|
495
|
+
official code). For decoupled-weight-decay optimizers such as
|
|
496
|
+
:class:`torch.optim.AdamW`, ``weight_decay`` is *not* equivalent to a Keras
|
|
497
|
+
``L2`` penalty.
|
|
498
|
+
|
|
499
|
+
Parameters
|
|
500
|
+
----------
|
|
501
|
+
conv_weight_decay : float
|
|
502
|
+
Weight decay applied to the convolution and TCN kernels. Defaults to
|
|
503
|
+
``0.009`` as in the official code.
|
|
504
|
+
dense_weight_decay : float
|
|
505
|
+
Weight decay applied to the final (dense) layer weights. Defaults to
|
|
506
|
+
``0.5`` as in the official code.
|
|
507
|
+
|
|
508
|
+
Returns
|
|
509
|
+
-------
|
|
510
|
+
list of dict
|
|
511
|
+
Parameter groups ready to be passed to a ``torch.optim`` optimizer:
|
|
512
|
+
convolution/TCN kernels, final dense weights, and everything else
|
|
513
|
+
(no decay).
|
|
514
|
+
|
|
515
|
+
Examples
|
|
516
|
+
--------
|
|
517
|
+
>>> model = ATCNet(n_chans=22, n_outputs=4, n_times=1125) # doctest: +SKIP
|
|
518
|
+
>>> opt = torch.optim.Adam(model.source_optimizer_param_groups()) # doctest: +SKIP
|
|
519
|
+
"""
|
|
520
|
+
|
|
521
|
+
def weight_leaf(module):
|
|
522
|
+
# Return the leaf weight Parameter, accounting for a registered
|
|
523
|
+
# max-norm parametrization (``conv_max_norm_const`` / MaxNormLinear).
|
|
524
|
+
if is_parametrized(module, "weight"):
|
|
525
|
+
return module.parametrizations.weight.original
|
|
526
|
+
return module.weight
|
|
527
|
+
|
|
528
|
+
conv_weights: list[nn.Parameter] = []
|
|
529
|
+
dense_weights: list[nn.Parameter] = []
|
|
530
|
+
decayed_ids: set[int] = set()
|
|
531
|
+
|
|
532
|
+
# Convolution / TCN kernels (Keras L2 = 0.009).
|
|
533
|
+
for parent in (self.conv_block, self.temporal_conv_nets):
|
|
534
|
+
for module in parent.modules():
|
|
535
|
+
if isinstance(module, (nn.Conv1d, nn.Conv2d)):
|
|
536
|
+
weight = weight_leaf(module)
|
|
537
|
+
conv_weights.append(weight)
|
|
538
|
+
decayed_ids.add(id(weight))
|
|
539
|
+
|
|
540
|
+
# Final dense / readout weights (Keras L2 = 0.5).
|
|
541
|
+
for module in self.final_layer.modules():
|
|
542
|
+
if isinstance(module, nn.Linear):
|
|
543
|
+
weight = weight_leaf(module)
|
|
544
|
+
dense_weights.append(weight)
|
|
545
|
+
decayed_ids.add(id(weight))
|
|
546
|
+
|
|
547
|
+
# Everything else (biases, BatchNorm, attention) is left undecayed.
|
|
548
|
+
other = [p for p in self.parameters() if id(p) not in decayed_ids]
|
|
549
|
+
|
|
550
|
+
return [
|
|
551
|
+
{"params": conv_weights, "weight_decay": conv_weight_decay},
|
|
552
|
+
{"params": dense_weights, "weight_decay": dense_weight_decay},
|
|
553
|
+
{"params": other, "weight_decay": 0.0},
|
|
554
|
+
]
|
|
555
|
+
|
|
556
|
+
|
|
557
|
+
def _register_max_norm(max_norm, *convs):
|
|
558
|
+
"""Register a max-norm weight parametrization on each given convolution.
|
|
559
|
+
|
|
560
|
+
No-op when ``max_norm`` is ``None`` or a passed layer is not a convolution
|
|
561
|
+
(e.g. an ``nn.Identity`` residual shortcut).
|
|
562
|
+
"""
|
|
563
|
+
if max_norm is None:
|
|
564
|
+
return
|
|
565
|
+
for conv in convs:
|
|
566
|
+
if isinstance(conv, (nn.Conv1d, nn.Conv2d)):
|
|
567
|
+
register_parametrization(conv, "weight", MaxNormParametrize(max_norm))
|
|
568
|
+
|
|
461
569
|
|
|
462
570
|
class _ConvBlock(nn.Module):
|
|
463
571
|
r"""Convolutional block proposed in ATCNet [1]_, inspired by the EEGNet.
|
|
@@ -487,6 +595,7 @@ class _ConvBlock(nn.Module):
|
|
|
487
595
|
pool_size_2=7,
|
|
488
596
|
depth_mult=2,
|
|
489
597
|
dropout=0.3,
|
|
598
|
+
conv_max_norm: Optional[float] = None,
|
|
490
599
|
):
|
|
491
600
|
super().__init__()
|
|
492
601
|
|
|
@@ -534,6 +643,9 @@ class _ConvBlock(nn.Module):
|
|
|
534
643
|
|
|
535
644
|
self.drop3 = nn.Dropout2d(dropout)
|
|
536
645
|
|
|
646
|
+
# Optional source-faithful max-norm constraint on the conv kernels.
|
|
647
|
+
_register_max_norm(conv_max_norm, self.conv1, self.conv2, self.conv3)
|
|
648
|
+
|
|
537
649
|
def forward(self, X):
|
|
538
650
|
# ----- Temporal convolution -----
|
|
539
651
|
# Dimension: (batch_size, 1, T, C)
|
|
@@ -662,6 +774,7 @@ class _TCNResidualBlock(nn.Module):
|
|
|
662
774
|
dropout=0.3,
|
|
663
775
|
activation: type[nn.Module] = nn.ELU,
|
|
664
776
|
dilation=1,
|
|
777
|
+
conv_max_norm: Optional[float] = None,
|
|
665
778
|
):
|
|
666
779
|
super().__init__()
|
|
667
780
|
self.activation = activation()
|
|
@@ -706,6 +819,10 @@ class _TCNResidualBlock(nn.Module):
|
|
|
706
819
|
else:
|
|
707
820
|
self.reshaping_conv = nn.Identity()
|
|
708
821
|
|
|
822
|
+
# Optional source-faithful max-norm constraint on the conv kernels
|
|
823
|
+
# (reshaping_conv may be nn.Identity, which the helper skips).
|
|
824
|
+
_register_max_norm(conv_max_norm, self.conv1, self.conv2, self.reshaping_conv)
|
|
825
|
+
|
|
709
826
|
def forward(self, X):
|
|
710
827
|
# Dimension: (batch_size, F2, Tw)
|
|
711
828
|
# ----- Double dilated convolutions -----
|
braindecode/models/attn_sleep.py
CHANGED
|
@@ -8,7 +8,6 @@ from copy import deepcopy
|
|
|
8
8
|
|
|
9
9
|
import torch
|
|
10
10
|
import torch.nn.functional as F
|
|
11
|
-
from mne.utils import deprecated
|
|
12
11
|
from torch import nn
|
|
13
12
|
|
|
14
13
|
from braindecode.models.base import EEGModuleMixin
|
|
@@ -538,12 +537,3 @@ class _PositionwiseFeedForward(nn.Module):
|
|
|
538
537
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
539
538
|
"""Implements FFN equation."""
|
|
540
539
|
return self.w_2(self.dropout(self.activate(self.w_1(x))))
|
|
541
|
-
|
|
542
|
-
|
|
543
|
-
@deprecated(
|
|
544
|
-
"`SleepStagerEldele2021` was renamed to `AttnSleep` in v1.12 to follow original author's name; this alias will be removed in v1.14."
|
|
545
|
-
)
|
|
546
|
-
class SleepStagerEldele2021(AttnSleep):
|
|
547
|
-
r"""Deprecated alias for SleepStagerEldele2021."""
|
|
548
|
-
|
|
549
|
-
pass
|