Hugging Face's logo

PCC-Lab-USI
/
UME-eda-foundation-model

Feature Extraction
PyTorch
Safetensors
electrodermal-activity
eda
wearables
physiological-signals
time-series
contrastive-learning
foundation-model
Model card Files
xet
 Community
UME-eda-foundation-model / README.md
bugcatcher's picture
bugcatcher
Update UME model files
dda346d verified
preview code
|
Raw History Blame Contribute Delete
21 kB
metadata 
license: mit
library_name: pytorch
pipeline_tag: feature-extraction
tags:
  - electrodermal-activity
  - eda
  - wearables
  - physiological-signals
  - time-series
  - contrastive-learning
  - foundation-model

🌸 UME: Foundation Model for Electrodermal Activity (EDA) Signals

Python 3.12+ PyTorch Lightning License: MIT DOI arXiv

See also 🫛 EDAMAME, the large-scale EDA pretraining collection we curated to train UME: repository · website (see Pretraining data).

🚀 Quick Use (UME) 

pip install git+https://github.com/LeonardoAlchieri/eda-foundation-models.git

import torch
from ume import UMEPipeline

# Downloads the base UME weights from the Hugging Face Hub
pipeline = UMEPipeline.from_pretrained()

# NOTE: default UME requires 60s windows (at 4Hz sampling rate) and 3 channels (mixed, phasic and tonic)
x = torch.randn(8, 240, 3)  # (batch, seq_len, channels)
embeddings = pipeline.embed(x)
print(embeddings.shape)  # (8, 64)

To load from a local clone instead, pass the weights file: UMEPipeline.from_pretrained(weights_path="model.safetensors") (see Checkpoints and Weights). A runnable walkthrough is in notebook/pipeline_demo.ipynb; for other ways to install, see Installation.

📋 Table of Contents

🎯 Overview

This repo contains the code and weights for UME, the first foundation model for wearable EDA data. UME is based on EfficientNet architecture and is trained using sample-wise contrastive learning, using the set of data augmentations from Katie Matton.

We release the weights for base UME, the default and suggested configuration, as well as other configurations (tiny, small and large). See below for more information on repo usage.

🪪 Model Card

  • What it is: an encoder that turns 60-second windows of wrist electrodermal activity (EDA) into 64-dimensional embeddings.
  • Architecture: EfficientNet-style 1D CNN (stem convolution, 16 MBConv1D blocks with squeeze-and-excitation, head convolution, global average pooling, linear projection). See ume/model/efficientnet_contrastive.py and UME Architecture.
  • Training objective: sample-wise contrastive learning (InfoNCE, temperature 0.1) between two augmented views of the same window, using the EDA augmentations of Matton et al. (2024).
  • Files: see Checkpoints and Weights.
  • Developed by: Leonardo Alchieri, Università della Svizzera italiana (USI).
  • Paper: doi.org/10.1145/3857981 · arXiv:2603.16878

Input

  • Shape (batch, 240, 3): 240 samples = 60 s at 4 Hz.
  • Channels, in order: mixed EDA, phasic component, tonic component.
  • Signals at other sampling rates or window lengths should be resampled and cut to match.

Intended use

  • Extracting features from wearable EDA for research on downstream tasks such as stress, cognitive load, or sleep/wake classification.
  • Fine-tuning on a labelled research dataset with UMEPipeline.fit (see Fine-Tune).

Out-of-scope use

  • Clinical diagnosis, or any medical or safety-critical decision about an individual.
  • Inferring sensitive traits or emotional states of people without their informed consent, including workplace or surveillance monitoring.
  • Signals that are not wrist EDA (e.g. PPG, ECG) or that are sampled very differently, without first validating performance.

Training data

Pretrained without labels on 17 public and research EDA datasets from the EDAMAME collection (see Pretraining data). The data is not distributed with the model. Recordings come mostly from research studies with wrist-worn devices such as the Empatica E4, which shapes the populations, devices and settings the model has seen.

Evaluation

Downstream evaluation (frozen features and a classical classifier, validated leave-one-person-out) is run by train.py through the optional edamame_downstream package. Results are reported in the paper.

Limitations and risks

  • Population and device bias. Performance can drop on populations, devices, body locations or recording conditions that are underrepresented in the pretraining data.
  • Artifacts. EDA is sensitive to motion, temperature and sensor contact. The model does not detect or flag poor-quality signal.
  • Individual differences. EDA varies a lot between people. Validate with person-independent splits (e.g. leave-one-person-out) rather than random splits, which overestimate performance.
  • Privacy. Physiological signals are personal data. Handle any data you use with the model according to its consent terms and applicable law.

🔧 Installation

Pixi (recommended)

The environment is defined in pixi.toml. Install Pixi, then:

git clone https://github.com/LeonardoAlchieri/eda-foundation-models.git
cd eda-foundation-models
pixi install --frozen

--frozen installs the default environment straight from pixi.lock. Without it, Pixi also checks the optional downstream environment, which needs the sibling repository described below.

pip 

pip install -e .            # inference and fine-tuning (UMEPipeline)
pip install -e ".[train]"   # + everything train.py needs for pretraining

Other extras: mantis (Mantis-8M backbone), dev (pytest, black, flake8, isort), notebook. See INSTALL.md.

Optional: downstream evaluation

After pretraining, train.py can evaluate the new model on labelled downstream tasks. This needs the separate edamame_downstream package, cloned next to this repository:

git clone https://github.com/LeonardoAlchieri/pretrained_foundation_models_physiological_data.git ../pretrained_foundation_models_physiological_data
pixi install -e downstream
pixi run -e downstream python train.py --config-name conf-ume

Without it, train.py pretrains as usual and skips the downstream step with a warning. Set downstream.enabled=false to skip it explicitly.

🎓 Fine-Tune (UME)

Beyond frozen-feature extraction, the pipeline can be supervised fine-tuned for classification through a scikit-learn-style API (fit / predict / predict_proba). A classification head (LayerNorm -> Linear by default, matching Mantis and Pulse-PPG) is attached on top of the encoder embedding and trained with cross-entropy.

Note: fine-tuning is currently supported only for the EfficientNet (UME) encoder (model_name="efficientnet"). MAE returns per-patch tokens rather than a pooled embedding, so it (and Mantis) raise NotImplementedError.

import torch
from ume import UMEPipeline

pipeline = UMEPipeline.from_pretrained()

X = torch.randn(256, 240, 3)              # (n_samples, seq_len, channels)
y = torch.randint(0, 2, (256,)).numpy()   # labels may be ints or strings

pipeline.fit(X, y, fine_tuning_type="full", num_epochs=20)

probs = pipeline.predict_proba(X)  # (n_samples, num_classes); column j -> pipeline.classes_[j]
preds = pipeline.predict(X)        # labels in the original space passed to `fit`

Labels are encoded internally to contiguous integer ids (sorted unique values, exposed as pipeline.classes_) and mapped back by predict.

Fine-tuning strategies

fine_tuning_type Behaviour
"full" (default) Fine-tune the whole encoder + head. Pass backbone_lr for a smaller encoder LR than the head's learning_rate.
"head" Freeze the encoder (linear probe) and train the head only; the backbone stays in eval mode so BatchNorm statistics are not updated.
"scratch" Re-initialise the encoder and train everything from random weights, ignoring the pretrained checkpoint — the "no pretraining" baseline.

Validation, early stopping and best-checkpoint selection

Supply a validation set either explicitly with validation_data=(X_val, y_val) or by holding out a stratified fraction of the training data with validation_split (e.g. 0.2). validation_data takes precedence when both are given.

pipeline.fit(
    X, y,
    fine_tuning_type="full",
    learning_rate=1e-3,
    backbone_lr=1e-4,
    num_epochs=100,
    validation_split=0.2,
    class_weight="balanced",      # inverse-frequency CE weights (or pass a tensor)
    early_stopping=True,
    monitor="val_macro_f1",       # or "val_acc" (maximized) / "val_loss" (minimized)
    patience=10,
    restore_best_weights=True,    # roll back to the best epoch when training ends
)

Without early_stopping, all num_epochs are run and no checkpoint selection happens. With it, training stops after patience epochs without a min_delta improvement in monitor, and restore_best_weights=True (the default) restores the encoder + head from the best epoch.

Refit on the full dataset

On small labelled datasets, permanently holding out labels for validation hurts the final model. refit_on_full_data=True (which requires early_stopping=True) instead runs training in two phases:

  1. Phase 1 trains with the validation hold-out purely to find the optimal training budget — the optimizer-step count of the best epoch (steps, not epochs, because the full dataset yields more batches per epoch).
  2. Phase 2 resets the encoder + head to their initial weights, rebuilds the optimizer, and retrains on 100% of the data (train + validation combined) for exactly that step budget, with no validation set.
pipeline.fit(
    X, y,
    validation_split=0.2,
    early_stopping=True,
    refit_on_full_data=True,
)

The refit model cannot be validated internally, so use an outer evaluation loop (e.g. leave-one-person-out) for unbiased performance estimates. restore_best_weights is ignored in this mode, since the refit replaces the phase-1 model entirely.

Key fit arguments

Argument Default Purpose
fine_tuning_type "full" "full" / "head" / "scratch" (see above).
learning_rate 1e-3 LR for the head (and the backbone when backbone_lr is None).
backbone_lr None Optional separate, typically smaller, encoder LR.
batch_size / num_epochs 64 / 100 Mini-batch size and epoch count.
weight_decay 0.0 AdamW weight decay.
validation_data / validation_split None / 0.0 Explicit validation set, or stratified hold-out fraction.
class_weight None "balanced" for inverse-frequency CE weights, or a (num_classes,) tensor.
early_stopping / patience / monitor / min_delta False / 10 / "val_macro_f1" / 0.0 Early-stopping controls.
restore_best_weights True Roll back to the best monitored epoch at the end of training.
refit_on_full_data False Two-phase budget-then-refit training (see above).
use_layer_norm True LayerNorm -> Linear head; set False for a bare linear head.
verbose True Print per-epoch loss / macro-F1.

🏋️ Retrain (UME)

We provide Hydra config files to reproduce the experiments from our work. Modify them to try different configurations. Retraining needs the pretraining data, which is not included in this repository (see Pretraining data).

Notes on the defaults:

  • Logging. Runs log to Weights & Biases under your default entity. Add pretrain/logger=csv (or tensorboard, none) to train without a W&B account.
  • Epoch length. An "epoch" is 3000 randomly sampled batches of 256 windows (pretrain.trainer.limit_train_batches), as for the released model. Set it to null to use the full training set per epoch.
  • Quick check. Add pretrain.trainer.fast_dev_run=5 to run a few batches end to end.

1. Train base UME 

pixi run python train.py --config-name conf-ume

2. Train other model and dataset sizes 

# Base (default)
python train.py --config-name conf-ume model_size=base

# Other presets
python train.py --config-name conf-ume model_size=tiny
python train.py --config-name conf-ume model_size=small
python train.py --config-name conf-ume model_size=large

# Base (default)
python train.py --config-name conf-ume pretrain.dataset.subsample=1.0

# Other presets
python train.py --config-name conf-ume pretrain.dataset.subsample=0.05
python train.py --config-name conf-ume pretrain.dataset.subsample=0.5
python train.py --config-name conf-ume pretrain.dataset.subsample=0.75

We also provide templates for sweeps for the model and dataset size (see sweep-ume-ablation-size.yaml and sweep-ume-ablation-subsample.yaml).

🫛 Pretraining data

UME was pretrained on 🫛 EDAMAME, which we release in two versions:

  • Open: freely downloadable.
  • Restricted: the full collection, available after signing a data sharing agreement.

See the EDAMAME repository for a description of the collection and the EDAMAME website for access to both versions. The data itself is not included in this repository.

The base UME was pretrained on these datasets from the collection (the list in configs/conf-ume.yaml):

BiHeartS, big-ideas-glycemic-wearable, Dynamics, EmpaticaE4Stress, EPM, LAUREATE, MEFAR, M2Sleep, PPG_FieldStudy, SEED, SEED-II-Lab, SEED-II-Wild, Stress Predict Dataset, ToadStool, WEEE, WESD, Workplaces.

train.py reads the data from pretrain.dataset.base_path (default data_online.nosync/), with one folder per dataset:

data_online.nosync/
└── <dataset_name>/
    └── *.h5          # one or more pandas HDFStore files

Each key in an HDF5 store is one recording, stored as a pandas DataFrame with the EDA channels as columns (mixed, phasic, tonic, sampled at 4 Hz). An optional Timestamp column is dropped. Each key's storer has an attrs.metadata dict with user, session, dataset_name and len_data (number of samples). The loader cuts 60 s windows (240 samples) on the fly with a stride of shift samples, and splits train and validation by recording, so windows from one recording never end up in both.

🏗️ UME Architecture (Base)

UME is an EfficientNet-inspired 1D encoder for contrastive learning:

Input (N, T, C)
  -> Stem Conv1D + BN + Swish
  -> MBConv1D blocks (repeated)
  -> Head Conv1D + Global Pool
  -> Linear embedding
  -> Contrastive loss on two augmented views

Base preset (default), set through the model_size presets in configs/conf-ume.yaml:

Parameter Base Value
stem_channels 64
mbconv_channels 64
num_mbconv_blocks 16
head_channels 248
embedding_dim 64
kernel_size 3
dropout 0.1
loss temperature 0.1
optimizer Adam, lr 5e-4, weight decay 1e-3

Contrastive training setup (augmentations, optimizer, scheduler):

🧪 UME Configurations

We recommend the base UME preset, and only its weights are released. The other presets and the pretraining-data subsamples used in our ablations can be retrained with the configs below.

UME model-size presets

Defined in configs/conf-ume.yaml:

Preset Stem MBConv Blocks Head Embedding Kernel Dropout
tiny 32 32 4 96 32 5 0.1
small 48 48 8 160 48 7 0.1
base 64 64 16 248 64 3 0.1
large 96 96 24 384 96 11 0.1

Sweeps over model size and pretraining-data size: sweep-ume-ablation-size.yaml and sweep-ume-ablation-subsample.yaml.

💾 Checkpoints and Weights

Three files hold the same base UME weights:

File Contents Use it for
model.safetensors + config.json Encoder weights in safetensors format, with the model config. Holds tensors only, so loading never runs code. This is what from_pretrained() downloads from the Hugging Face Hub. Inference and fine-tuning (recommended).
artifacts/ume_base_inference.pt Encoder state_dict and model config in one PyTorch file. Loads with torch.load(..., weights_only=True). Same as above, without the safetensors dependency.
artifacts/ume_base_training.ckpt Full PyTorch Lightning checkpoint, including optimizer and scheduler state. Needs weights_only=False. Resuming or inspecting training.

Are they the same model? Yes. All three hold the same 366 encoder tensors, bit-for-bit, and give identical embeddings. They differ only in what else they contain:

  • ume_base_training.ckpt is about 3× larger because it also stores the training state (mostly the Adam optimizer state, plus epoch, step, callback and hyperparameter records). Its weight keys carry a model. prefix, which from_pretrained strips.
  • ume_base_training.ckpt has no model config: from_pretrained infers the architecture from the tensor shapes and gets the same model. The other two files include the config.
  • Only ume_base_training.ckpt needs weights_only=False, which can run code stored in the file, so load it only from a source you trust.

Use ume_base_training.ckpt only to resume pretraining. For inference and fine-tuning, use model.safetensors or ume_base_inference.pt.

UMEPipeline.from_pretrained accepts any of them through weights_path. With no weights_path, it downloads model.safetensors and config.json from the Hub repo given by repo_id (default: the released base UME); revision and cache_dir are passed to the download.

📚 Citation

If you use UME, please cite the paper (doi.org/10.1145/3857981):

@article{alchieri2026foundation,
  author     = {Alchieri, Leonardo and Garzon, Matteo and Alecci, Lidia and Bombassei De Bona, Francesco and Gjoreski, Martin and De Felice, Giovanni and Santini, Silvia},
  title      = {A Foundation Model for Electrodermal Activity Data},
  journal    = {Proc. ACM Interact. Mob. Wearable Ubiquitous Technol.},
  volume     = {10},
  number     = {4},
  articleno  = {202},
  numpages   = {40},
  year       = {2026},
  month      = dec,
  publisher  = {Association for Computing Machinery},
  address    = {New York, NY, USA},
  doi        = {10.1145/3857981},
  url        = {https://doi.org/10.1145/3857981}
}

📄 License

This project is licensed under the MIT License. See LICENSE.

People Centered Computing Lab @ USI, Switzerland.