Training Pipeline#

This example shows the complete training pipeline from raw data to model training. Each step is explained and demonstrated.

from pathlib import Path

import torch

# Get the path to the data file
# Find data directory relative to myoverse package (works in all contexts)
import myoverse
_pkg_dir = Path(myoverse.__file__).parent.parent
DATA_DIR = _pkg_dir / "examples" / "data"
if not DATA_DIR.exists():
    DATA_DIR = Path.cwd() / "examples" / "data"

# Determine device
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

# Step 1: Create Dataset with Preprocessing
# ------------------------------------------
# Use DatasetCreator with Modality transforms for pre-storage processing.
# Here we use the EMBC paper configuration as an example.

from myoverse.datasets import DatasetCreator, Modality, embc_kinematics_transform

print("=" * 60)
print("STEP 1: Dataset Creation")
print("=" * 60)

creator = DatasetCreator(
    modalities={
        # EMG: raw continuous data (320 channels from 5 electrode grids)
        "emg": Modality(
            path=DATA_DIR / "emg.pkl",
            dims=("channel", "time"),
        ),
        # Kinematics: apply transform to flatten and remove wrist
        # (21, 3, time) -> (60, time)
        "kinematics": Modality(
            path=DATA_DIR / "kinematics.pkl",
            dims=("dof", "time"),
            transform=embc_kinematics_transform(),
        ),
    },
    sampling_frequency=2048.0,
    tasks_to_use=["1", "2"],
    save_path=DATA_DIR / "tutorial_dataset.zip",
    test_ratio=0.2,
    val_ratio=0.2,
    debug_level=1,
)
creator.create()
Using device: cpu
============================================================
STEP 1: Dataset Creation
============================================================
────────────────────────── STARTING DATASET CREATION ───────────────────────────

                             Dataset Configuration
┏━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Parameter               ┃ Value                                              ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ Modalities              │ emg, kinematics                                    │
│ Sampling frequency (Hz) │ 2048.0                                             │
│ Save path               │ /home/runner/work/MyoVerse/MyoVerse/examples/data… │
│ Test ratio              │ 0.2                                                │
│ Validation ratio        │ 0.2                                                │
└─────────────────────────┴────────────────────────────────────────────────────┘

Processing 2 tasks: 1, 2

Dataset Structure
├── emg dims=('channel', 'time')
│   ├── Task 1: (320, 20440)
│   └── Task 2: (320, 20440)
└── kinematics dims=('dof', 'time')
    ├── Task 1: (60, 20440)
    └── Task 2: (60, 20440)

─────────────────────────────── PROCESSING TASKS ───────────────────────────────

  Processing task 2 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:00
────────────────────────── DATASET CREATION COMPLETED ──────────────────────────

               Dataset Summary
 Split       emg              kinematics
 training    1: (320, 16352)  1: (60, 16352)
             2: (320, 16352)  2: (60, 16352)
 validation  1: (320, 816)    1: (60, 816)
             2: (320, 816)    2: (60, 816)
 testing     1: (320, 3272)   1: (60, 3272)
             2: (320, 3272)   2: (60, 3272)

Total size: 26.09 MB
─────────────────── Dataset Creation Successfully Completed! ───────────────────

Step 2: Define Training Transforms#

Transforms are applied on-the-fly during training. - embc_train_transform: Creates dual representation (raw + lowpass) + augmentation - embc_eval_transform: Same processing without augmentation

from myoverse.datasets import embc_eval_transform, embc_target_transform, embc_train_transform

print()
print("=" * 60)
print("STEP 2: Training Transforms")
print("=" * 60)

# Training: dual representation + noise augmentation
train_tf = embc_train_transform(augmentation="noise")
print(f"Train transform: {train_tf}")

# Validation: dual representation only (no augmentation)
val_tf = embc_eval_transform()
print(f"Val transform: {val_tf}")

# Target: average kinematics over window -> single prediction per DOF
target_tf = embc_target_transform()
print(f"Target transform: {target_tf}")
============================================================
STEP 2: Training Transforms
============================================================
Train transform: Compose(
    Stack(dim='time', transforms={'raw': Identity(dim='time'), 'filtered': Lowpass(dim='time', cutoff=20.0, fs=2048.0, order=4, Q=0.707)}, stack_dim='representation')
    GaussianNoise(dim='time', std=0.1, p=1.0)
)
Val transform: Compose(
    Stack(dim='time', transforms={'raw': Identity(dim='time'), 'filtered': Lowpass(dim='time', cutoff=20.0, fs=2048.0, order=4, Q=0.707)}, stack_dim='representation')
)
Target transform: Mean(dim='time', keepdim=False)

Step 3: Create DataModule#

DataModule handles: - Loading from zarr directly to tensors (GPU if available) - On-the-fly windowing (no pre-chunking needed) - Input/target selection (decided at training time, not storage time) - Transform application - Batching and DataLoader creation

from myoverse.datasets import DataModule

print()
print("=" * 60)
print("STEP 3: DataModule Setup")
print("=" * 60)

dm = DataModule(
    data_path=DATA_DIR / "tutorial_dataset.zip",
    # Select which modalities are inputs vs targets
    inputs=["emg"],
    targets=["kinematics"],
    # Windowing parameters
    window_size=192,  # ~94ms at 2048Hz
    window_stride=64,  # For val/test (deterministic sliding window)
    n_windows_per_epoch=500,  # For training (random positions) - small for demo
    # Transforms (applied on-the-fly)
    train_transform=train_tf,
    val_transform=val_tf,
    target_transform=target_tf,  # Average kinematics over window
    # DataLoader settings
    batch_size=32,
    num_workers=0,  # Set to 4+ for parallel loading
    # Device: load directly to GPU if available
    device=DEVICE,
)

# Setup creates the datasets
dm.setup("fit")

print(f"Training samples per epoch: {len(dm.train_dataloader()) * dm.batch_size}")
print(f"Validation batches: {len(dm.val_dataloader())}")
============================================================
STEP 3: DataModule Setup
============================================================
Training samples per epoch: 512
Validation batches: 1

Step 4: Inspect Batch Structure#

With single input/target, DataModule returns tensors directly (for compatibility with existing models). With multiple inputs/targets, it returns dicts.

print()
print("=" * 60)
print("STEP 4: Batch Structure")
print("=" * 60)

batch = next(iter(dm.train_dataloader()))
emg_batch, kin_batch = batch

print(f"EMG input shape: {emg_batch.shape}")
print(f"EMG input device: {emg_batch.device}")
print(f"Kinematics target shape: {kin_batch.shape}")

# EMG shape explanation (with Stack transform):
# - Batch size: 32
# - Representations: 2 (raw, filtered)
# - Channels: varies based on data
# - Time: 192 samples
print()
print("EMG shape = (batch, representation, channel, time) via Stack transform")
============================================================
STEP 4: Batch Structure
============================================================
EMG input shape: torch.Size([32, 2, 320, 192])
EMG input device: cpu
Kinematics target shape: torch.Size([32, 60])

EMG shape = (batch, representation, channel, time) via Stack transform

Step 5: Create Model#

RaulNetV16 expects: - Input: (batch, 2, channels, time) - 2 representations - Output: (batch, 60) - 60 DOF predictions

print()
print("=" * 60)
print("STEP 5: Model Setup")
print("=" * 60)

from myoverse.models import RaulNetV16

# Get actual channel count from data
n_channels = emg_batch.shape[2]
n_grids = n_channels // 64 if n_channels >= 64 else 1

model = RaulNetV16(
    learning_rate=1e-4,
    nr_of_input_channels=2,  # raw + filtered
    input_length__samples=192,
    nr_of_outputs=60,  # 60 DOF
    nr_of_electrode_grids=n_grids,
    nr_of_electrodes_per_grid=64,
    cnn_encoder_channels=(4, 1, 1),
    mlp_encoder_channels=(8, 8),
    event_search_kernel_length=31,
    event_search_kernel_stride=8,
)

# Move model to same device as data
model = model.to(DEVICE)

print(f"Model: {model.__class__.__name__}")
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
============================================================
STEP 5: Model Setup
============================================================
Model: RaulNetV16
Parameters: 14,720

Step 6: Forward Pass Test#

Named tensors are used during transforms for dimension-awareness, but stripped in the collate function before passing to the model.

print()
print("=" * 60)
print("STEP 6: Forward Pass Test")
print("=" * 60)

print(f"DataModule output: {emg_batch.shape}")
print("(Names stripped in collate - ready for model)")

# Quick forward pass test
model.eval()
with torch.no_grad():
    output = model(emg_batch)
print(f"Model output: {output.shape}")
============================================================
STEP 6: Forward Pass Test
============================================================
DataModule output: torch.Size([32, 2, 320, 192])
(Names stripped in collate - ready for model)
Model output: torch.Size([32, 60])

Step 7: Training Loop#

Use PyTorch Lightning Trainer for training. Note: For real training, increase max_epochs and n_windows_per_epoch.

import lightning as L

print()
print("=" * 60)
print("STEP 7: Training (1 epoch, 500 windows)")
print("=" * 60)

torch.set_float32_matmul_precision("medium")  # For performance on some CPUs

trainer = L.Trainer(
    accelerator="auto",
    devices=1,
    precision="32",  # Use 32-bit for CPU compatibility
    max_epochs=1,
    log_every_n_steps=5,
    logger=False,
    enable_checkpointing=False,
    enable_progress_bar=True,
)

# Train for 1 epoch
trainer.fit(model, datamodule=dm)
============================================================
STEP 7: Training (1 epoch, 500 windows)
============================================================
/home/runner/work/MyoVerse/MyoVerse/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=3` in the `DataLoader` to improve performance.
┏━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━┳━━━━━━━┓
┃   ┃ Name        ┃ Type       ┃ Params ┃ Mode ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━╇━━━━━━━┩
│ 0 │ criterion   │ L1Loss     │      0 │ eval │     0 │
│ 1 │ cnn_encoder │ Sequential │ 11.7 K │ eval │     0 │
│ 2 │ mlp         │ Sequential │  3.0 K │ eval │     0 │
└───┴─────────────┴────────────┴────────┴──────┴───────┘
Trainable params: 14.7 K
Non-trainable params: 0
Total params: 14.7 K
Total estimated model params size (MB): 0
Modules in train mode: 0
Modules in eval mode: 19
Total FLOPs: 0
/home/runner/work/MyoVerse/MyoVerse/.venv/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=3` in the `DataLoader` to improve performance.
/home/runner/work/MyoVerse/MyoVerse/.venv/lib/python3.12/site-packages/lightning/pytorch/core/module.py:522: You called `self.log('val/val_loss', ..., logger=True)` but have no logger configured. You can enable one by doing `Trainer(logger=ALogger(...))`
/home/runner/work/MyoVerse/MyoVerse/.venv/lib/python3.12/site-packages/lightning/pytorch/loops/fit_loop.py:534: Found 20 module(s) in eval mode at the start of training. This may lead to unexpected behavior during training. If this is intentional, you can ignore this warning.
/home/runner/work/MyoVerse/MyoVerse/.venv/lib/python3.12/site-packages/lightning/pytorch/core/module.py:522: You called `self.log('train/loss', ..., logger=True)` but have no logger configured. You can enable one by doing `Trainer(logger=ALogger(...))`
Epoch 0/0  ━━━━━━━━━━━━━━━━━ 16/16 0:00:09 • 0:00:00 1.69it/s loss_step: 54.514
                                                              val_loss: 54.312
                                                              loss_epoch: 54.414

Summary#

The complete pipeline:

  1. DatasetCreator - Store continuous data with pre-processing transforms

  2. Modality.transform - Pre-storage transforms (e.g., flatten kinematics)

  3. DataModule - Load directly to GPU, window, select inputs/targets

  4. train_transform - On-the-fly transforms (filtering, augmentation)

  5. Model - Your neural network

Key benefits: - Modular: swap transforms without changing dataset - Efficient: zarr -> GPU loading (kvikio if available) - Flexible: input/target selection at training time - Named tensors: dimension-aware transforms

print()
print("=" * 60)
print("PIPELINE COMPLETE")
print("=" * 60)
============================================================
PIPELINE COMPLETE
============================================================

Total running time of the script: (0 minutes 14.063 seconds)

Estimated memory usage: 678 MB

Gallery generated by Sphinx-Gallery