ProtoSleepNet and ProtoSleepTransformer

Two pretrained sleep stagers, trained on 28 sites from 12 public and clinical datasets and published on the Hugging Face Hub in 4rooms/physioex:

Hub name	Epoch encoder	Sequence length L	Parameters
protosleepnet-gagliardi (PSN)	SeqSleepNet	20	0.79 M
protosleeptransformer-gagliardi (PST)	SleepTransformer	21	3.98 M

Both are physioex.models.protosleepnet.ProtoSleepNet networks: a per-channel epoch encoder, a channel mixer that learns how much to trust each modality, a residual sequence encoder and a linear classifier. They are the models of the paper Prototype-based interpretable sleep staging with physiologically meaningful sub-stage pattern discovery (npj Digital Medicine; code and analyses in guidogagl/protosleepnet), retrained on many sites at once. The weights of the paper itself (one dataset per model) are a separate archive.

Quick start

import torch
from physioex.models import load_from_pretrained

model = load_from_pretrained("protosleeptransformer-gagliardi", verbose=True)   # eval mode, CPU

# (batch, L, channels, T, F): log-power STFT spectrograms of EEG, EOG, EMG (T=29, F=129)
x = torch.randn(2, 21, 3, 29, 129)
with torch.no_grad():
    logits = model(x)            # (2, 21, 5): W, N1, N2, N3, REM

verbose=True prints the per-site test metrics stored with the weights. The input is what the seqsleepnet pipeline preset produces for the channels ["EEG", "EOG", "EMG"], so the models can be fed directly from any physioex dataset:

from physioex.data import get_dataset

ds = get_dataset("hmc")(channels=["EEG", "EOG", "EMG"], pipelines="seqsleepnet", sequence_length=21)

encode(x) returns the contextualised epoch embeddings (B, L, d_model). forward(x, quantize=True) replaces the epoch embeddings by their nearest prototype, but it needs a codebook (set_codebook) that is learned after training from the epoch embeddings; the checkpoints on the Hub contain no codebook, and quantize=False (the default) is the plain classifier. The prototype analyses of the paper are in its repository, and relevance maps for these models are in Explainability (LRP supports the default, non-quantised forward).

Training data

Both models were trained jointly on the training subjects of 28 sites. The test and validation subjects of every site are disjoint from the training subjects (split by subject, per site).

Dataset	Sites	Source
Sleep-EDF	sleepedf	PhysioNet
HMC	hmc	PhysioNet
DCSM	dcsm	ERDA / University of Copenhagen
MASS	mass_ss01 … mass_ss05 (cohorts SS1-SS5)	CEAMS
SHHS	shhs	NSRR
MESA	mesa	NSRR
MrOS	mros	NSRR
HomePAP	hpap	NSRR
WSC	wsc	NSRR
STAGES	stages_BOGN, stages_GSBB, stages_GSDV, stages_GSLH, stages_GSSA, stages_GSSW, stages_MSMI, stages_MSNF, stages_MSQW, stages_MSTH, stages_MSTR, stages_STLK, stages_STNF (13 clinical sites)	NSRR
Alzheimer’s disease study, healthy controls	alzheimers_hc	UZ Leuven
Parkinson’s disease study, healthy older adults, night recordings	parkinsons_night_hoa	UZ Leuven

Recipe (the same for both models): 3 channels (EEG, EOG, EMG), seqsleepnet pipeline, site-balanced sampling (make_site_sampler, every site gets the same probability mass, see Data), batch size 32, learning rate 1e-4 with a reduce-on-plateau scheduler, no weight decay, 300,000 windows per epoch, at most 20 epochs, channel dropout 0.5 in the channel mixer, seed 42. The released checkpoint is the one with the highest macro-averaged validation kappa among the evaluated checkpoints (up to 30 validation subjects per site): step 161,500 (epoch 17) for PSN and step 129,000 (epoch 13) for PST. Neighbouring checkpoints differ by about 0.001 kappa, so the choice is not critical.

Performance by training site

Metrics are computed on the test subjects of each site, pooling the epochs of all of them (this is not a mean over nights), with full-night sliding-window voting (windows of L epochs at L different offsets, overlapping predictions averaged). Accuracy, Cohen’s kappa and macro F1 over the five AASM stages. Read the tables with these limits in mind:

• The sites were seen in training (other subjects of the same sites), so this is not external validation: a new site, montage or population can perform differently.

• Single training run (one seed), no confidence intervals. Several sites have 2-6 test subjects, so their values are noisy; SHHS is by far the largest test set (842 nights).

• Per-site values mix population, scoring protocol and recording hardware; do not read a difference between two sites as a difference in model quality without checking the labels.

The two single-centre studies at the end of the training table (alzheimers_hc, parkinsons_night_hoa) were part of training but their per-site metrics are not published, so the tables below list 26 of the 28 sites.

ProtoSleepNet (protosleepnet-gagliardi)

Site	Test subjects	Accuracy	Cohen’s kappa	Macro F1
dcsm	26	0.867	0.809	0.774
hmc	15	0.814	0.754	0.784
hpap	25	0.847	0.790	0.801
mass_ss01	5	0.811	0.724	0.761
mass_ss02	2	0.910	0.858	0.837
mass_ss03	6	0.888	0.835	0.844
mass_ss04	4	0.871	0.814	0.834
mass_ss05	3	0.891	0.841	0.821
mesa	206	0.842	0.776	0.771
mros	219	0.859	0.795	0.757
shhs	842	0.872	0.818	0.780
sleepedf	15	0.842	0.778	0.794
stages_BOGN	8	0.865	0.800	0.806
stages_GSBB	3	0.814	0.729	0.691
stages_GSDV	19	0.804	0.679	0.697
stages_GSLH	3	0.812	0.732	0.740
stages_GSSA	5	0.826	0.696	0.673
stages_GSSW	8	0.836	0.722	0.711
stages_MSMI	6	0.840	0.764	0.787
stages_MSNF	4	0.854	0.735	0.742
stages_MSQW	15	0.805	0.711	0.734
stages_MSTH	3	0.823	0.745	0.771
stages_MSTR	29	0.849	0.779	0.800
stages_STLK	16	0.802	0.708	0.704
stages_STNF	49	0.691	0.591	0.630
wsc	235	0.833	0.735	0.729

ProtoSleepTransformer (protosleeptransformer-gagliardi)

Site	Test subjects	Accuracy	Cohen’s kappa	Macro F1
dcsm	26	0.870	0.814	0.781
hmc	15	0.815	0.754	0.789
hpap	25	0.857	0.803	0.812
mass_ss01	5	0.810	0.723	0.766
mass_ss02	2	0.899	0.841	0.833
mass_ss03	6	0.894	0.843	0.846
mass_ss04	4	0.870	0.812	0.829
mass_ss05	3	0.888	0.838	0.824
mesa	206	0.842	0.778	0.778
mros	219	0.872	0.811	0.765
shhs	842	0.880	0.829	0.797
sleepedf	15	0.842	0.775	0.780
stages_BOGN	8	0.873	0.811	0.817
stages_GSBB	3	0.828	0.749	0.708
stages_GSDV	19	0.777	0.649	0.688
stages_GSLH	3	0.790	0.699	0.698
stages_GSSA	5	0.821	0.692	0.676
stages_GSSW	8	0.812	0.675	0.674
stages_MSMI	6	0.834	0.763	0.794
stages_MSNF	4	0.853	0.743	0.757
stages_MSQW	15	0.819	0.730	0.749
stages_MSTH	3	0.823	0.743	0.763
stages_MSTR	29	0.844	0.773	0.796
stages_STLK	16	0.812	0.724	0.717
stages_STNF	49	0.720	0.632	0.672
wsc	235	0.845	0.750	0.739

The same numbers are stored with the weights (metrics.json in each model folder on the Hub).

Intended use and limits

For research on sleep staging and its interpretability; not a medical device, and not for clinical diagnosis. The models expect the spectrogram input described above: a different preprocessing, signal amplitude scale or montage changes the input distribution, so check the staging on a sample of your own data before relying on it. The channel set is fixed to EEG, EOG and EMG (the channel mixer was trained with random channel dropout, which makes the models tolerant to a missing channel, but this was not evaluated per site here).

Reproducibility

Every model folder on the Hub has config.json with the factory call and the training summary, including the SHA-256 of model.pt:

Hub name	Checkpoint step	SHA-256 of model.pt
protosleepnet-gagliardi	161,500	1156a3fe18d3bc8a83a775bcc8e1a1ba08c0c4624b3415a4fc07982e907136db
protosleeptransformer-gagliardi	129,000	08d0fbc806fc08dc7d0167e1c2b81b0e2b4f5495bad2c1e5dda6bde14972e603

physioex/examples/pretrained/protosleepnet/export_release.py builds the folder from a training checkpoint and its per-site evaluation; nothing is recomputed, and the metrics are copied unchanged.

API

class physioex.models.protosleepnet.ProtoSleepNet(epoch_encoder, sequence_encoder, classifier, n_channels=1, n_classes=5, d_model=128, use_channel_mixer=False, cdropout=0.5, cm_n_heads=4, cm_d_ff=256, cm_n_layers=1)

Bases: Module

Template-based sleep staging model with dual residual connections.

Wraps any epoch_encoder + sequence_encoder with: - Optional channel mixer (for multi-channel input) - Residual sequence encoder: z = h + seq(h) - Single classifier

Parameters:

• epoch_encoder (Module) – maps (N, 1, T, F) → (N, d_model) for per-channel, or (N, C, T, F) → (N, d_model) for single-channel

• sequence_encoder (Module) – maps (B, L, d_model) → (B, L, d_model)

• classifier (Module) – maps (N, d_model) → (N, n_classes)

• n_channels (int) – number of input channels (1 = no mixer)

• n_classes (int) – number of output classes

• d_model (int) – epoch encoder output dimension

• use_channel_mixer (bool) – enable channel mixer (requires n_channels > 1)

• cdropout (float) – channel dropout probability

• cm_n_heads (int) – mixer transformer heads

• cm_d_ff (int) – mixer transformer feedforward dim

• cm_n_layers (int) – mixer transformer layers

encode(x, quantize=False)

Encode input spectrograms to contextualized per-epoch embeddings.

Calls epoch_encode → sequence_encode.

Parameters:

• x – (B, L, C, T, F) spectrogram input.

• quantize – if True, replace epoch embeddings with nearest codebook entry. Requires set_codebook() first.

Returns:

(B, L, d_model) contextualized epoch embeddings.

epoch_encode(x, quantize=False)

Encode input spectrograms to per-epoch embeddings.

Pipeline: per-channel epoch encoding → [channel mixer] → [VQ] → deep supervision (epoch_logits stored in metrics).

Parameters:

• x – (B, L, C, T, F) spectrogram input.

• quantize – if True, replace epoch embeddings with nearest codebook entry. Requires set_codebook() first.

Returns:

(B, L, d_model) epoch embeddings (pre-sequence).

forward(x, quantize=False)

Forward pass: encode → classify.

Parameters:

• x – (B, L, C, T, F) spectrogram input.

• quantize – if True, replace epoch embeddings with nearest codebook entry. Requires set_codebook() first.

Returns:

(B, L, n_classes) per-epoch logits.

classmethod from_seq_sleep_net(n_channels=1, n_classes=5, F=129, D=32, nfft=256, lowfreq=0, highfreq=50, fs=100, seqnhidden1=64, seqnlayer1=4, attentionsize=32, seqnhidden2=64, seqnlayer2=4, **mixer_kwargs)

Build ProtoSleepNet with SeqSleepNet backbone.

Parameters:

• n_channels (int)

• n_classes (int)

• F (int)

• D (int)

• nfft (int)

• lowfreq (int)

• highfreq (int)

• fs (int)

• seqnhidden1 (int)

• seqnlayer1 (int)

• attentionsize (int)

• seqnhidden2 (int)

• seqnlayer2 (int)

Return type:

ProtoSleepNet

classmethod from_sleep_transformer(n_channels=1, n_classes=5, d_model=128, n_heads=8, n_epoch_layers=4, n_seq_layers=4, d_ff=1024, d_clf=1024, dropout=0.1, attention_size=128, **mixer_kwargs)

Build ProtoSleepNet with SleepTransformer backbone.

Parameters:

• n_channels (int)

• n_classes (int)

• d_model (int)

• n_heads (int)

• n_epoch_layers (int)

• n_seq_layers (int)

• d_ff (int)

• d_clf (int)

• dropout (float)

• attention_size (int)

Return type:

ProtoSleepNet

sequence_encode(h)

Contextualize epoch embeddings with residual sequence encoding.

Parameters:

h – (B, L, d_model) epoch embeddings.

Returns:

h + seq(h).

Return type:

(B, L, d_model) contextualized embeddings

set_codebook(codebook)

Set VQ codebook for prototype quantization.

Parameters:

codebook – (M, d_model) numpy array or torch tensor.