Skip to content

deepecgkit

Top-level exports from the DeepECG-Kit package.

deepecgkit

DeepECG-Kit: Deep learning library for ECG analysis and arrhythmia classification.

ECGDataModule

Data module for ECG datasets.

Handles dataset creation, train/val/test splitting, and DataLoader construction.

Source code in deepecgkit/datasets/modules.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
class ECGDataModule:
    """Data module for ECG datasets.

    Handles dataset creation, train/val/test splitting, and DataLoader construction.
    """

    def __init__(
        self,
        dataset: BaseECGDataset | Dataset | None = None,
        dataset_class: type[BaseECGDataset] | None = None,
        data_dir: str | None = None,
        sampling_rate: int | None = None,
        leads: list | None = None,
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        batch_size: int = 32,
        val_split: float = 0.2,
        test_split: float = 0.1,
        num_workers: int = 0,
        seed: int = 42,
        stratify: bool = True,
        pin_memory: bool = False,
        persistent_workers: bool = True,
        prefetch_factor: int = 2,
        verbose: bool = True,
        dataset_kwargs: dict | None = None,
    ):
        """Initialize the ECG data module.

        Args:
            dataset: PyTorch dataset instance (optional)
            dataset_class: Class of dataset to create (optional)
            data_dir: Directory containing ECG data files (optional, uses dataset's default if None)
            sampling_rate: Rate the dataset resamples its signals to (Hz). None keeps the
                dataset class's own default, which is its native or recommended rate;
                after setup() the attribute holds the rate actually in effect
            leads: List of leads to use (e.g., ['I', 'II', 'III'] for standard leads)
            transform: Optional transform to be applied to the ECG signals
            target_transform: Optional transform to be applied to the labels
            download: Whether to download the dataset if it doesn't exist
            batch_size: Batch size for data loaders
            val_split: Fraction of data to use for validation
            test_split: Fraction of data to use for testing
            num_workers: Number of worker processes for data loading. The built-in
                datasets hold their signals in memory, so the default of 0 is the
                fastest setting; with more, guard the calling script with
                ``if __name__ == "__main__":`` on macOS and Windows
            seed: Random seed for reproducibility
            stratify: Whether to use stratified splitting based on labels
            dataset_kwargs: Additional keyword arguments to pass to the dataset class
        """
        self.dataset = dataset
        self.dataset_class = dataset_class
        self.data_dir = data_dir
        self.sampling_rate = sampling_rate
        self.leads = leads
        self.transform = transform
        self.target_transform = target_transform
        self.download = download
        self.batch_size = batch_size
        self.val_split = val_split
        self.test_split = test_split
        self.num_workers = num_workers
        self.seed = seed
        self.stratify = stratify
        self.pin_memory = pin_memory
        self.persistent_workers = persistent_workers
        self.prefetch_factor = prefetch_factor
        self.verbose = verbose
        self.dataset_kwargs = dataset_kwargs or {}

        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None
        self._split_settings: tuple | None = None

    def setup(self, stage: str | None = None) -> None:
        """Set up the datasets.

        The dataset is loaded and split on the first call. Later calls, which
        the trainer makes for every stage, return at once unless the dataset or
        a split setting has changed since.

        Args:
            stage: Current stage ('fit', 'validate', 'test', or 'predict')
        """
        if self.dataset is None:
            if self.dataset_class is None:
                raise ValueError("dataset_class must be provided if dataset is None")

            if self.data_dir is None:
                self.data_dir = self.dataset_class.get_default_data_dir()

            init_kwargs = {
                "data_dir": self.data_dir,
                "transform": self.transform,
                "target_transform": self.target_transform,
                "download": self.download,
            }

            # Only an explicit rate overrides the dataset's own default: a
            # blanket value here would silently resample every dataset.
            if self.sampling_rate is not None:
                init_kwargs["sampling_rate"] = self.sampling_rate

            if self.leads is not None:
                init_kwargs["leads"] = self.leads

            init_kwargs.update(self.dataset_kwargs)

            self.dataset = self.dataset_class(**init_kwargs)

        effective_rate = getattr(self.dataset, "sampling_rate", None)
        if effective_rate is not None:
            self.sampling_rate = effective_rate

        split_settings = (
            id(self.dataset),
            self.val_split,
            self.test_split,
            self.seed,
            self.stratify,
        )
        if self.train_dataset is not None and split_settings == self._split_settings:
            return
        self._split_settings = split_settings

        stratify_labels = None
        if (
            self.stratify and len(self.dataset) >= 8
        ):  # Need at least 2 samples per class for 4 classes
            try:
                stacked = self._collect_labels()

                # For multi-label (2D), convert rows to string keys for stratification
                if stacked.ndim > 1:
                    stratify_labels = np.array(
                        ["_".join(str(int(v)) for v in row) for row in stacked]
                    )
                else:
                    stratify_labels = stacked

                # Check if we have enough samples per class for stratification
                _, counts = np.unique(stratify_labels, return_counts=True)
                if np.min(counts) < 2:
                    warnings.warn(
                        "Stratified splitting disabled: at least one class has fewer "
                        "than 2 samples.",
                        RuntimeWarning,
                        stacklevel=2,
                    )
                    stratify_labels = None
            except (RuntimeError, ValueError, TypeError, IndexError) as error:
                warnings.warn(
                    f"Stratified splitting disabled: could not collect labels ({error}).",
                    RuntimeWarning,
                    stacklevel=2,
                )
                stratify_labels = None

        # Extract groups for patient-level splitting (prevents data leakage)
        groups = None
        group_ids = getattr(self.dataset, "group_ids", None) or getattr(
            self.dataset, "record_names", None
        )
        if group_ids:
            if len(group_ids) != len(self.dataset):
                raise ValueError(
                    f"{type(self.dataset).__name__} exposes {len(group_ids)} group ids "
                    f"for {len(self.dataset)} samples. Grouped splitting cannot be "
                    "performed safely."
                )
            groups = np.array(group_ids)
            if self.verbose:
                n_groups = len(np.unique(groups))
                logger.info(f"Using grouped splitting ({n_groups} groups)")
        else:
            warnings.warn(
                f"{type(self.dataset).__name__} exposes no group ids, so the split is "
                "performed at sample level. If samples are segments of longer recordings, "
                "segments of the same recording will appear in both train and test and the "
                "reported scores will be optimistically biased.",
                RuntimeWarning,
                stacklevel=2,
            )

        splitter = DataSplitter(
            dataset=self.dataset,
            val_split=self.val_split,
            test_split=self.test_split,
            seed=self.seed,
            stratify=stratify_labels,
            groups=groups,
        )
        self.train_dataset, self.val_dataset, self.test_dataset = splitter.split()

        if self.verbose:
            logger.info(f"Training set size: {len(self.train_dataset)}")
            logger.info(f"Validation set size: {len(self.val_dataset)}")
            logger.info(f"Test set size: {len(self.test_dataset)}")

    def _collect_labels(self) -> np.ndarray:
        """Collect every sample's label, for stratifying the split.

        A dataset that can hand its labels over directly is asked first;
        otherwise every sample is read for its label.
        """
        size = len(self.dataset)
        direct = getattr(self.dataset, "stratification_labels", None)
        labels: np.ndarray | None = direct() if callable(direct) else None
        if labels is None or len(labels) != size:
            labels = torch.stack([self.dataset[i][1] for i in range(size)]).numpy()
        return labels

    def _loader(self, dataset: Any, shuffle: bool) -> DataLoader:
        """Build a loader over one split.

        Raises:
            RuntimeError: If setup() has not been called.
        """
        if dataset is None:
            raise RuntimeError("Call setup() before using the data module")
        options: dict[str, Any] = {}
        if self.num_workers > 0:
            # These two only mean something with worker processes, and PyTorch
            # before 2.0 rejects a prefetch_factor given without any.
            options["persistent_workers"] = self.persistent_workers
            options["prefetch_factor"] = self.prefetch_factor
        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            **options,
        )

    def train_dataloader(self) -> DataLoader:
        """Get the training data loader."""
        return self._loader(self.train_dataset, shuffle=True)

    def val_dataloader(self) -> DataLoader:
        """Get the validation data loader."""
        return self._loader(self.val_dataset, shuffle=False)

    def test_dataloader(self) -> DataLoader:
        """Get the test data loader."""
        return self._loader(self.test_dataset, shuffle=False)

    def describe(self) -> dict[str, Any]:
        """Describe the data a model is trained on, in the form checkpoints record.

        ECGTrainer stores this as its data_config, which is what lets a
        checkpoint be used on new recordings: it names the preprocessing the
        inputs need and the classes the outputs stand for.

        Returns:
            Mapping with the registered dataset name, its preprocessing
            (sampling rate, segment length, normalisation), the class names,
            the sample shape, the labelling mode, and the split protocol.
        """
        if self.dataset is None:
            raise RuntimeError("Call setup() before describing the data module")
        dataset = self.dataset

        input_channels = input_length = None
        if len(dataset) > 0:
            sample = dataset[0][0]
            if getattr(sample, "ndim", 0) >= 2:
                input_channels, input_length = int(sample.shape[0]), int(sample.shape[-1])

        get_preprocessing = getattr(dataset, "get_preprocessing", None)
        class_names = getattr(dataset, "class_names", None)
        return {
            "dataset": getattr(type(dataset), "_registry_name", None),
            "preprocessing": get_preprocessing() if callable(get_preprocessing) else None,
            "class_names": list(class_names) if class_names else None,
            "input_channels": input_channels,
            "input_length": input_length,
            "multi_label": bool(getattr(dataset, "multi_label", False)),
            "binary_classification": bool(getattr(dataset, "binary_classification", False)),
            "split": {
                "seed": self.seed,
                "val_split": self.val_split,
                "test_split": self.test_split,
            },
        }

    def get_metadata(self) -> dict:
        """Get metadata about the dataset.

        Returns:
            Dictionary containing dataset metadata
        """
        if self.dataset is None:
            raise RuntimeError("Call setup() before getting metadata")
        if isinstance(self.dataset, BaseECGDataset):
            return self.dataset.get_metadata()
        return {
            "dataset_size": len(self.dataset),
            "batch_size": self.batch_size,
            "num_workers": self.num_workers,
        }

    def print_metadata(self) -> None:
        """Print metadata about the dataset."""
        print(self.get_metadata())

setup 

setup(stage: str | None = None) -> None

Set up the datasets.

The dataset is loaded and split on the first call. Later calls, which the trainer makes for every stage, return at once unless the dataset or a split setting has changed since.

Parameters:

Name Type Description Default
stage str | None

Current stage ('fit', 'validate', 'test', or 'predict')

 None
Source code in deepecgkit/datasets/modules.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def setup(self, stage: str | None = None) -> None:
    """Set up the datasets.

    The dataset is loaded and split on the first call. Later calls, which
    the trainer makes for every stage, return at once unless the dataset or
    a split setting has changed since.

    Args:
        stage: Current stage ('fit', 'validate', 'test', or 'predict')
    """
    if self.dataset is None:
        if self.dataset_class is None:
            raise ValueError("dataset_class must be provided if dataset is None")

        if self.data_dir is None:
            self.data_dir = self.dataset_class.get_default_data_dir()

        init_kwargs = {
            "data_dir": self.data_dir,
            "transform": self.transform,
            "target_transform": self.target_transform,
            "download": self.download,
        }

        # Only an explicit rate overrides the dataset's own default: a
        # blanket value here would silently resample every dataset.
        if self.sampling_rate is not None:
            init_kwargs["sampling_rate"] = self.sampling_rate

        if self.leads is not None:
            init_kwargs["leads"] = self.leads

        init_kwargs.update(self.dataset_kwargs)

        self.dataset = self.dataset_class(**init_kwargs)

    effective_rate = getattr(self.dataset, "sampling_rate", None)
    if effective_rate is not None:
        self.sampling_rate = effective_rate

    split_settings = (
        id(self.dataset),
        self.val_split,
        self.test_split,
        self.seed,
        self.stratify,
    )
    if self.train_dataset is not None and split_settings == self._split_settings:
        return
    self._split_settings = split_settings

    stratify_labels = None
    if (
        self.stratify and len(self.dataset) >= 8
    ):  # Need at least 2 samples per class for 4 classes
        try:
            stacked = self._collect_labels()

            # For multi-label (2D), convert rows to string keys for stratification
            if stacked.ndim > 1:
                stratify_labels = np.array(
                    ["_".join(str(int(v)) for v in row) for row in stacked]
                )
            else:
                stratify_labels = stacked

            # Check if we have enough samples per class for stratification
            _, counts = np.unique(stratify_labels, return_counts=True)
            if np.min(counts) < 2:
                warnings.warn(
                    "Stratified splitting disabled: at least one class has fewer "
                    "than 2 samples.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                stratify_labels = None
        except (RuntimeError, ValueError, TypeError, IndexError) as error:
            warnings.warn(
                f"Stratified splitting disabled: could not collect labels ({error}).",
                RuntimeWarning,
                stacklevel=2,
            )
            stratify_labels = None

    # Extract groups for patient-level splitting (prevents data leakage)
    groups = None
    group_ids = getattr(self.dataset, "group_ids", None) or getattr(
        self.dataset, "record_names", None
    )
    if group_ids:
        if len(group_ids) != len(self.dataset):
            raise ValueError(
                f"{type(self.dataset).__name__} exposes {len(group_ids)} group ids "
                f"for {len(self.dataset)} samples. Grouped splitting cannot be "
                "performed safely."
            )
        groups = np.array(group_ids)
        if self.verbose:
            n_groups = len(np.unique(groups))
            logger.info(f"Using grouped splitting ({n_groups} groups)")
    else:
        warnings.warn(
            f"{type(self.dataset).__name__} exposes no group ids, so the split is "
            "performed at sample level. If samples are segments of longer recordings, "
            "segments of the same recording will appear in both train and test and the "
            "reported scores will be optimistically biased.",
            RuntimeWarning,
            stacklevel=2,
        )

    splitter = DataSplitter(
        dataset=self.dataset,
        val_split=self.val_split,
        test_split=self.test_split,
        seed=self.seed,
        stratify=stratify_labels,
        groups=groups,
    )
    self.train_dataset, self.val_dataset, self.test_dataset = splitter.split()

    if self.verbose:
        logger.info(f"Training set size: {len(self.train_dataset)}")
        logger.info(f"Validation set size: {len(self.val_dataset)}")
        logger.info(f"Test set size: {len(self.test_dataset)}")

train_dataloader 

train_dataloader() -> DataLoader

Get the training data loader.

Source code in deepecgkit/datasets/modules.py 
258
259
260
def train_dataloader(self) -> DataLoader:
    """Get the training data loader."""
    return self._loader(self.train_dataset, shuffle=True)

val_dataloader 

val_dataloader() -> DataLoader

Get the validation data loader.

Source code in deepecgkit/datasets/modules.py 
262
263
264
def val_dataloader(self) -> DataLoader:
    """Get the validation data loader."""
    return self._loader(self.val_dataset, shuffle=False)

test_dataloader 

test_dataloader() -> DataLoader

Get the test data loader.

Source code in deepecgkit/datasets/modules.py 
266
267
268
def test_dataloader(self) -> DataLoader:
    """Get the test data loader."""
    return self._loader(self.test_dataset, shuffle=False)

describe 

describe() -> dict[str, Any]

Describe the data a model is trained on, in the form checkpoints record.

ECGTrainer stores this as its data_config, which is what lets a checkpoint be used on new recordings: it names the preprocessing the inputs need and the classes the outputs stand for.

Returns:

Type Description
dict[str, Any]

Mapping with the registered dataset name, its preprocessing
dict[str, Any]

(sampling rate, segment length, normalisation), the class names,
dict[str, Any]

the sample shape, the labelling mode, and the split protocol.
Source code in deepecgkit/datasets/modules.py 
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
def describe(self) -> dict[str, Any]:
    """Describe the data a model is trained on, in the form checkpoints record.

    ECGTrainer stores this as its data_config, which is what lets a
    checkpoint be used on new recordings: it names the preprocessing the
    inputs need and the classes the outputs stand for.

    Returns:
        Mapping with the registered dataset name, its preprocessing
        (sampling rate, segment length, normalisation), the class names,
        the sample shape, the labelling mode, and the split protocol.
    """
    if self.dataset is None:
        raise RuntimeError("Call setup() before describing the data module")
    dataset = self.dataset

    input_channels = input_length = None
    if len(dataset) > 0:
        sample = dataset[0][0]
        if getattr(sample, "ndim", 0) >= 2:
            input_channels, input_length = int(sample.shape[0]), int(sample.shape[-1])

    get_preprocessing = getattr(dataset, "get_preprocessing", None)
    class_names = getattr(dataset, "class_names", None)
    return {
        "dataset": getattr(type(dataset), "_registry_name", None),
        "preprocessing": get_preprocessing() if callable(get_preprocessing) else None,
        "class_names": list(class_names) if class_names else None,
        "input_channels": input_channels,
        "input_length": input_length,
        "multi_label": bool(getattr(dataset, "multi_label", False)),
        "binary_classification": bool(getattr(dataset, "binary_classification", False)),
        "split": {
            "seed": self.seed,
            "val_split": self.val_split,
            "test_split": self.test_split,
        },
    }

get_metadata 

get_metadata() -> dict

Get metadata about the dataset.

Returns:

Type Description
dict

Dictionary containing dataset metadata
Source code in deepecgkit/datasets/modules.py 
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
def get_metadata(self) -> dict:
    """Get metadata about the dataset.

    Returns:
        Dictionary containing dataset metadata
    """
    if self.dataset is None:
        raise RuntimeError("Call setup() before getting metadata")
    if isinstance(self.dataset, BaseECGDataset):
        return self.dataset.get_metadata()
    return {
        "dataset_size": len(self.dataset),
        "batch_size": self.batch_size,
        "num_workers": self.num_workers,
    }

print_metadata 

print_metadata() -> None

Print metadata about the dataset.

Source code in deepecgkit/datasets/modules.py 
325
326
327
def print_metadata(self) -> None:
    """Print metadata about the dataset."""
    print(self.get_metadata())

ECGTrainer

Trainer for ECG signal classification and regression models.

Wraps a plain nn.Module and provides fit/test methods with built-in early stopping, checkpointing, LR scheduling, and CSV metric logging.

Parameters:

Name Type Description Default
model Module

The ECG model to train (any nn.Module)

 required
train_config dict[str, Any] | None

Optional dictionary of training settings; every key has a default, so ECGTrainer(model) trains a multi-class classifier: - learning_rate: Learning rate for optimizer (default 1e-3) - scheduler: Dict with 'factor' and 'patience' for ReduceLROnPlateau (default factor 0.5, patience 5) - binary_classification: Bool, if True uses BCE loss for binary tasks - multi_label: Bool, if True uses BCE loss for multi-label tasks - task_type: 'classification' (default) or 'regression' - pos_weight: Optional list of positive class weights for BCE loss

 None
device str | device

Device to train on ('auto', 'cpu', 'cuda', 'mps')

 'auto'
use_plateau_scheduler bool

If True, uses ReduceLROnPlateau, else StepLR

 True
Example

model = KanResWideX(input_channels=1, output_size=4) trainer = ECGTrainer(model) trainer.fit(data_module, epochs=50)

Source code in deepecgkit/training/train.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
class ECGTrainer:
    """
    Trainer for ECG signal classification and regression models.

    Wraps a plain nn.Module and provides fit/test methods with built-in
    early stopping, checkpointing, LR scheduling, and CSV metric logging.

    Args:
        model: The ECG model to train (any nn.Module)
        train_config: Optional dictionary of training settings; every key has
            a default, so ``ECGTrainer(model)`` trains a multi-class classifier:
            - learning_rate: Learning rate for optimizer (default 1e-3)
            - scheduler: Dict with 'factor' and 'patience' for
              ReduceLROnPlateau (default factor 0.5, patience 5)
            - binary_classification: Bool, if True uses BCE loss for binary tasks
            - multi_label: Bool, if True uses BCE loss for multi-label tasks
            - task_type: 'classification' (default) or 'regression'
            - pos_weight: Optional list of positive class weights for BCE loss
        device: Device to train on ('auto', 'cpu', 'cuda', 'mps')
        use_plateau_scheduler: If True, uses ReduceLROnPlateau, else StepLR

    Example:
        >>> model = KanResWideX(input_channels=1, output_size=4)
        >>> trainer = ECGTrainer(model)
        >>> trainer.fit(data_module, epochs=50)
    """

    def __init__(
        self,
        model: nn.Module,
        train_config: dict[str, Any] | None = None,
        device: str | torch.device = "auto",
        use_plateau_scheduler: bool = True,
    ) -> None:
        self.model = model
        train_config = dict(train_config or {})
        train_config.setdefault("learning_rate", DEFAULT_LEARNING_RATE)
        train_config["scheduler"] = {**DEFAULT_SCHEDULER, **(train_config.get("scheduler") or {})}
        self.train_config = train_config

        self.learning_rate = train_config["learning_rate"]
        self.scheduler_factor = train_config["scheduler"]["factor"]
        self.scheduler_patience = train_config["scheduler"]["patience"]
        self.use_plateau_scheduler = use_plateau_scheduler
        self.multi_label = train_config.get("multi_label", False)

        if self.multi_label or train_config.get("binary_classification", False):
            pos_weight = train_config.get("pos_weight")
            if pos_weight is not None:
                pos_weight = torch.tensor(pos_weight, dtype=torch.float32)
            self.criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
        elif train_config.get("task_type", "classification") == "classification":
            self.criterion = torch.nn.CrossEntropyLoss()
        else:
            self.criterion = torch.nn.MSELoss()

        self.device = resolve_device(device)

        self.model.to(self.device)
        self.criterion.to(self.device)

        self.optimizer: torch.optim.Optimizer | None = None
        self.scheduler: Any = None
        self.test_predictions: list[torch.Tensor] = []
        self.test_targets: list[torch.Tensor] = []
        self.test_probabilities: list[torch.Tensor] = []
        self.test_logits: list[torch.Tensor] = []
        self.log_dir: str | None = None
        self.best_checkpoint_path: str | None = None
        self.best_val_loss = float("inf")
        #: How to rebuild the model: registry name, lead count, output size.
        self.model_config: dict[str, Any] = {}
        #: What the model was trained on: dataset, preprocessing, class names.
        #: Both are written into every checkpoint and filled in by fit() when
        #: left empty, so a checkpoint can be loaded and used on its own.
        self.data_config: dict[str, Any] = {}
        self._gradient_clip_val: float | None = None

    @property
    def _is_binary(self) -> bool:
        return isinstance(self.criterion, torch.nn.BCEWithLogitsLoss) and not self.multi_label

    @property
    def _is_classification(self) -> bool:
        return isinstance(self.criterion, torch.nn.CrossEntropyLoss | torch.nn.BCEWithLogitsLoss)

    def _calculate_loss(self, y_hat: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        if self.multi_label:
            return self.criterion(y_hat, y.float())
        if self._is_binary:
            return self.criterion(y_hat.squeeze(-1), y.float())
        if isinstance(self.criterion, torch.nn.CrossEntropyLoss):
            return self.criterion(y_hat, y.long())
        return self.criterion(y_hat.float(), y.float())

    def _compute_acc(self, y_hat: torch.Tensor, y: torch.Tensor) -> float:
        with torch.no_grad():
            if self.multi_label:
                preds = (y_hat > 0).float()
                return (preds == y.float()).float().mean().item()
            if self._is_binary:
                preds = (y_hat.squeeze(-1) > 0).long()
                return (preds == y.long()).float().mean().item()
            return (torch.argmax(y_hat, dim=1) == y.long()).float().mean().item()

    def _setup_optimizer(self) -> None:
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
        if self.use_plateau_scheduler:
            self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer,
                mode="min",
                factor=self.scheduler_factor,
                patience=self.scheduler_patience,
            )
        else:
            self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=1, gamma=0.9)

    def _run_epoch(self, dataloader: DataLoader, train: bool = True) -> tuple[float, float | None]:
        if train:
            self.model.train()
        else:
            self.model.eval()

        # Sums over samples, not batches: the last batch is usually smaller, and
        # counting it as a full one would give its samples more than their share.
        total_loss = 0.0
        total_acc = 0.0
        n_samples = 0

        ctx = torch.no_grad() if not train else contextlib.nullcontext()
        with ctx:
            for batch in dataloader:
                x, y = batch
                x, y = x.to(self.device), y.to(self.device)

                y_hat = self.model(x)
                loss = self._calculate_loss(y_hat, y)

                if train:
                    if hasattr(self.model, "l2_regularization_loss"):
                        loss = loss + self.model.l2_regularization_loss()
                    self.optimizer.zero_grad()
                    loss.backward()
                    if self._gradient_clip_val is not None:
                        torch.nn.utils.clip_grad_norm_(
                            self.model.parameters(), self._gradient_clip_val
                        )
                    self.optimizer.step()

                batch_size = len(y)
                total_loss += loss.item() * batch_size
                if self._is_classification:
                    total_acc += self._compute_acc(y_hat, y) * batch_size
                n_samples += batch_size
                if hasattr(dataloader, "set_postfix"):
                    dataloader.set_postfix(loss=f"{total_loss / n_samples:.4f}", refresh=False)

        avg_loss = total_loss / max(n_samples, 1)
        avg_acc = total_acc / max(n_samples, 1) if self._is_classification else None
        return avg_loss, avg_acc

    def fit(
        self,
        data_module: Any,
        epochs: int = 50,
        early_stopping_patience: int = 10,
        checkpoint_dir: str | Path | None = None,
        log_dir: str | Path | None = None,
        progress_bar: bool = True,
        gradient_clip_val: float | None = None,
        save_top_k: int = 3,
        restore_best: bool = True,
    ) -> None:
        """Train the model.

        On completion the weights achieving the lowest validation loss are restored,
        so any subsequent test() or validate() call evaluates the selected model
        rather than the final epoch. Requires checkpoint_dir to be set; without it
        no checkpoint exists to restore and the final-epoch weights are kept.

        Args:
            data_module: ECGDataModule (or any object with train_dataloader/val_dataloader)
            epochs: Maximum number of training epochs
            early_stopping_patience: Stop after this many epochs without val_loss improvement
            checkpoint_dir: Directory to save checkpoints (None to disable)
            log_dir: Directory to save CSV metrics log (None to disable)
            progress_bar: Whether to show tqdm progress: one bar over the epochs and,
                on an interactive terminal, a second one over the batches of the
                current epoch
            gradient_clip_val: Max gradient norm for clipping (None to disable)
            save_top_k: Number of best checkpoints to keep
            restore_best: Reload the best-validation-loss checkpoint when training ends
        """
        self._gradient_clip_val = gradient_clip_val
        self._setup_optimizer()

        if hasattr(data_module, "setup"):
            data_module.setup(stage="fit")

        train_loader = data_module.train_dataloader()
        val_loader = data_module.val_dataloader()
        self._describe_run(data_module)

        if checkpoint_dir is not None:
            checkpoint_dir = Path(checkpoint_dir)
            checkpoint_dir.mkdir(parents=True, exist_ok=True)

        csv_writer = None
        csv_file = None
        if log_dir is not None:
            self.log_dir = str(log_dir)
            os.makedirs(log_dir, exist_ok=True)
            metrics_path = Path(log_dir) / "metrics.csv"
            csv_file = open(metrics_path, "w", newline="")
            fieldnames = ["epoch", "train_loss", "val_loss"]
            if self._is_classification:
                fieldnames.extend(["train_acc", "val_acc"])
            csv_writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
            csv_writer.writeheader()

        self.best_val_loss = float("inf")
        patience_counter = 0
        saved_checkpoints = []

        use_tqdm = progress_bar and tqdm is not None
        epoch_iter = tqdm(range(epochs), desc="Training") if use_tqdm else range(epochs)
        # A bar per batch is only readable where it can redraw in place; in a
        # log file it would add a line every refresh.
        show_batches = use_tqdm and sys.stderr.isatty()

        try:
            for epoch in epoch_iter:
                batches: Any = train_loader
                if show_batches:
                    batches = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}", leave=False)
                train_loss, train_acc = self._run_epoch(batches, train=True)
                val_loss, val_acc = self._run_epoch(val_loader, train=False)

                if self.use_plateau_scheduler:
                    self.scheduler.step(val_loss)
                else:
                    self.scheduler.step()

                if use_tqdm:
                    desc = f"Epoch {epoch + 1}/{epochs} | train_loss={train_loss:.4f} val_loss={val_loss:.4f}"
                    if self._is_classification:
                        desc += f" val_acc={val_acc:.4f}"
                    epoch_iter.set_description(desc)

                if csv_writer is not None:
                    row = {
                        "epoch": epoch + 1,
                        "train_loss": f"{train_loss:.6f}",
                        "val_loss": f"{val_loss:.6f}",
                    }
                    if self._is_classification:
                        row["train_acc"] = f"{train_acc:.6f}"
                        row["val_acc"] = f"{val_acc:.6f}"
                    csv_writer.writerow(row)
                    csv_file.flush()

                if val_loss < self.best_val_loss:
                    self.best_val_loss = val_loss
                    patience_counter = 0

                    if checkpoint_dir is not None:
                        ckpt_path = (
                            checkpoint_dir / f"epoch={epoch + 1:02d}-val_loss={val_loss:.4f}.pt"
                        )
                        self.save_checkpoint(str(ckpt_path), epoch=epoch + 1)
                        self.best_checkpoint_path = str(ckpt_path)
                        saved_checkpoints.append((val_loss, str(ckpt_path)))
                        saved_checkpoints.sort(key=lambda x: x[0])
                        while len(saved_checkpoints) > save_top_k:
                            _, old_path = saved_checkpoints.pop()
                            if os.path.exists(old_path):
                                os.remove(old_path)
                else:
                    patience_counter += 1
                    if patience_counter >= early_stopping_patience:
                        if not use_tqdm:
                            logger.info(f"Early stopping at epoch {epoch + 1}")
                        break
        finally:
            if csv_file is not None:
                csv_file.close()

        if restore_best:
            self.restore_best_checkpoint()

    def _describe_run(self, data_module: Any) -> None:
        """Fill in whatever model_config and data_config leave unset.

        Callers such as the CLI set both explicitly; this covers direct use of
        the trainer, so its checkpoints still name their architecture and
        record the preprocessing their inputs went through.

        Args:
            data_module: The data module being trained on
        """
        described: dict[str, Any] = {}
        describe = getattr(data_module, "describe", None)
        if callable(describe):
            try:
                described = describe() or {}
            except Exception as error:  # a description must never stop a training run
                logger.debug(f"Could not describe the data module: {error}")

        if not self.data_config and described:
            self.data_config = described

        if not self.model_config.get("name"):
            name = getattr(type(self.model), "_registry_name", None)
            if name is not None:
                self.model_config["name"] = name
        if self.model_config.get("input_channels") is None and described.get("input_channels"):
            self.model_config["input_channels"] = described["input_channels"]
        if self.model_config.get("output_size") is None:
            from deepecgkit.models.factory import infer_output_size

            output_size = infer_output_size(
                self.model.state_dict(), getattr(self.model, "HEAD_ATTR", None)
            )
            if output_size is not None:
                self.model_config["output_size"] = output_size

    def restore_best_checkpoint(self) -> bool:
        """Reload the weights of the best-validation-loss checkpoint written during fit.

        Returns:
            True if a checkpoint was restored, False if none was available.
        """
        if self.best_checkpoint_path is None or not os.path.exists(self.best_checkpoint_path):
            if self.best_checkpoint_path is not None:
                warnings.warn(
                    f"Best checkpoint {self.best_checkpoint_path} is missing; "
                    "evaluation will use final-epoch weights.",
                    RuntimeWarning,
                    stacklevel=2,
                )
            return False

        checkpoint = load_checkpoint_file(self.best_checkpoint_path)
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model.to(self.device)
        return True

    def test(self, data_module: Any) -> dict[str, float]:
        """Evaluate the model on the test set.

        Args:
            data_module: ECGDataModule (or any object with test_dataloader)

        Returns:
            Dict with test_loss and test_acc (if classification)
        """
        if hasattr(data_module, "setup"):
            data_module.setup(stage="test")

        test_loader = data_module.test_dataloader()
        return self._evaluate_loader(test_loader)

    def validate(self, data_module: Any) -> dict[str, float]:
        """Evaluate the model on the validation set.

        Args:
            data_module: ECGDataModule (or any object with val_dataloader)

        Returns:
            Dict with val_loss and val_acc (if classification)
        """
        if hasattr(data_module, "setup"):
            data_module.setup(stage="validate")

        val_loader = data_module.val_dataloader()
        return self._evaluate_loader(val_loader, prefix="val")

    def _evaluate_loader(self, dataloader: DataLoader, prefix: str = "test") -> dict[str, float]:
        self.model.eval()
        self.test_predictions = []
        self.test_targets = []
        self.test_probabilities = []
        self.test_logits = []

        total_loss = 0.0
        total_acc = 0.0
        n_samples = 0

        with torch.no_grad():
            for batch in dataloader:
                x, y = batch
                x, y = x.to(self.device), y.to(self.device)
                y_hat = self.model(x)
                loss = self._calculate_loss(y_hat, y)
                batch_size = len(y)
                total_loss += loss.item() * batch_size
                self.test_logits.append(y_hat.detach().cpu())

                if self._is_classification:
                    if self.multi_label:
                        probs = torch.sigmoid(y_hat)
                        preds = (probs > 0.5).long()
                        acc = (preds == y.long()).float().mean().item()
                        self.test_predictions.append(preds.cpu())
                        self.test_targets.append(y.long().cpu())
                        self.test_probabilities.append(probs.cpu())
                    elif self._is_binary:
                        probs_pos = torch.sigmoid(y_hat.squeeze(-1))
                        probs = torch.stack([1 - probs_pos, probs_pos], dim=1)
                        preds = (probs_pos > 0.5).long()
                        acc = (preds == y.long()).float().mean().item()
                        self.test_predictions.append(preds.cpu())
                        self.test_targets.append(y.long().cpu())
                        self.test_probabilities.append(probs.cpu())
                    else:
                        probs = torch.softmax(y_hat, dim=1)
                        preds = torch.argmax(probs, dim=1)
                        acc = (preds == y.long()).float().mean().item()
                        self.test_predictions.append(preds.cpu())
                        self.test_targets.append(y.long().cpu())
                        self.test_probabilities.append(probs.cpu())
                    total_acc += acc * batch_size

                n_samples += batch_size

        avg_loss = total_loss / max(n_samples, 1)
        results = {f"{prefix}_loss": avg_loss}
        if self._is_classification:
            results[f"{prefix}_acc"] = total_acc / max(n_samples, 1)
        return results

    def get_test_results(self) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None]:
        """Get test predictions, targets, and probabilities as numpy arrays.

        Returns:
            Tuple of (predictions, targets, probabilities) as numpy arrays,
            or (None, None, None) if no test results available.
        """
        if not self.test_predictions:
            return None, None, None
        return (
            torch.cat(self.test_predictions).numpy(),
            torch.cat(self.test_targets).numpy(),
            torch.cat(self.test_probabilities).numpy(),
        )

    def get_test_logits(self) -> np.ndarray | None:
        """Get the raw model outputs from the most recent evaluation pass.

        Logits are kept alongside probabilities because temperature scaling
        needs them and the softmax and sigmoid that produce the probabilities
        are not losslessly invertible in float32: a probability saturates to
        exactly 1.0 once its logit passes about 16.6, which is well inside the
        range an overconfident model reaches.

        Returns:
            Raw outputs as a numpy array, shaped (n_samples, output_size), or
            None when no evaluation has been run.
        """
        if not self.test_logits:
            return None
        return torch.cat(self.test_logits).numpy()

    def save_checkpoint(self, path: str | Path, epoch: int | None = None) -> None:
        """Save a checkpoint.

        Args:
            path: File path to save to
            epoch: Current epoch number (optional)
        """
        checkpoint: dict[str, Any] = {
            "model_state_dict": self.model.state_dict(),
            "train_config": self.train_config,
            "model_config": self.model_config,
            "data_config": self.data_config,
            "epoch": epoch,
            "best_val_loss": self.best_val_loss,
            "deepecgkit_version": _package_version(),
        }
        if self.optimizer is not None:
            checkpoint["optimizer_state_dict"] = self.optimizer.state_dict()
        torch.save(checkpoint, path)

    @classmethod
    def load_checkpoint(
        cls,
        path: str | Path,
        model: nn.Module | None = None,
        device: str | torch.device = "auto",
    ) -> "ECGTrainer":
        """Load a trainer from a checkpoint.

        Args:
            path: Path to checkpoint file
            model: Model instance to load weights into. When omitted, the
                architecture recorded in the checkpoint is rebuilt.
            device: Device to load onto

        Returns:
            ECGTrainer instance with loaded weights

        Raises:
            ValueError: If no model is given and the checkpoint does not
                record which architecture it holds.
        """
        checkpoint = load_checkpoint_file(path)

        if model is None:
            from deepecgkit.models.factory import build_model_from_checkpoint

            model, _ = build_model_from_checkpoint(path)
        else:
            model.load_state_dict(checkpoint["model_state_dict"])

        trainer = cls(
            model=model,
            train_config=checkpoint.get("train_config"),
            device=device,
        )
        trainer.model_config = dict(checkpoint.get("model_config") or {})
        trainer.data_config = dict(checkpoint.get("data_config") or {})
        trainer.best_val_loss = checkpoint.get("best_val_loss", float("inf"))
        return trainer

    @staticmethod
    def seed_everything(seed: int) -> None:
        """Set random seeds for reproducibility.

        Args:
            seed: Random seed value
        """
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)

fit 

fit(
    data_module: Any,
    epochs: int = 50,
    early_stopping_patience: int = 10,
    checkpoint_dir: str | Path | None = None,
    log_dir: str | Path | None = None,
    progress_bar: bool = True,
    gradient_clip_val: float | None = None,
    save_top_k: int = 3,
    restore_best: bool = True,
) -> None

Train the model.

On completion the weights achieving the lowest validation loss are restored, so any subsequent test() or validate() call evaluates the selected model rather than the final epoch. Requires checkpoint_dir to be set; without it no checkpoint exists to restore and the final-epoch weights are kept.

Parameters:

Name Type Description Default
data_module Any

ECGDataModule (or any object with train_dataloader/val_dataloader)

 required
epochs int

Maximum number of training epochs

 50
early_stopping_patience int

Stop after this many epochs without val_loss improvement

 10
checkpoint_dir str | Path | None

Directory to save checkpoints (None to disable)

 None
log_dir str | Path | None

Directory to save CSV metrics log (None to disable)

 None
progress_bar bool

Whether to show tqdm progress: one bar over the epochs and, on an interactive terminal, a second one over the batches of the current epoch

 True
gradient_clip_val float | None

Max gradient norm for clipping (None to disable)

 None
save_top_k int

Number of best checkpoints to keep

 3
restore_best bool

Reload the best-validation-loss checkpoint when training ends

 True
Source code in deepecgkit/training/train.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
def fit(
    self,
    data_module: Any,
    epochs: int = 50,
    early_stopping_patience: int = 10,
    checkpoint_dir: str | Path | None = None,
    log_dir: str | Path | None = None,
    progress_bar: bool = True,
    gradient_clip_val: float | None = None,
    save_top_k: int = 3,
    restore_best: bool = True,
) -> None:
    """Train the model.

    On completion the weights achieving the lowest validation loss are restored,
    so any subsequent test() or validate() call evaluates the selected model
    rather than the final epoch. Requires checkpoint_dir to be set; without it
    no checkpoint exists to restore and the final-epoch weights are kept.

    Args:
        data_module: ECGDataModule (or any object with train_dataloader/val_dataloader)
        epochs: Maximum number of training epochs
        early_stopping_patience: Stop after this many epochs without val_loss improvement
        checkpoint_dir: Directory to save checkpoints (None to disable)
        log_dir: Directory to save CSV metrics log (None to disable)
        progress_bar: Whether to show tqdm progress: one bar over the epochs and,
            on an interactive terminal, a second one over the batches of the
            current epoch
        gradient_clip_val: Max gradient norm for clipping (None to disable)
        save_top_k: Number of best checkpoints to keep
        restore_best: Reload the best-validation-loss checkpoint when training ends
    """
    self._gradient_clip_val = gradient_clip_val
    self._setup_optimizer()

    if hasattr(data_module, "setup"):
        data_module.setup(stage="fit")

    train_loader = data_module.train_dataloader()
    val_loader = data_module.val_dataloader()
    self._describe_run(data_module)

    if checkpoint_dir is not None:
        checkpoint_dir = Path(checkpoint_dir)
        checkpoint_dir.mkdir(parents=True, exist_ok=True)

    csv_writer = None
    csv_file = None
    if log_dir is not None:
        self.log_dir = str(log_dir)
        os.makedirs(log_dir, exist_ok=True)
        metrics_path = Path(log_dir) / "metrics.csv"
        csv_file = open(metrics_path, "w", newline="")
        fieldnames = ["epoch", "train_loss", "val_loss"]
        if self._is_classification:
            fieldnames.extend(["train_acc", "val_acc"])
        csv_writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        csv_writer.writeheader()

    self.best_val_loss = float("inf")
    patience_counter = 0
    saved_checkpoints = []

    use_tqdm = progress_bar and tqdm is not None
    epoch_iter = tqdm(range(epochs), desc="Training") if use_tqdm else range(epochs)
    # A bar per batch is only readable where it can redraw in place; in a
    # log file it would add a line every refresh.
    show_batches = use_tqdm and sys.stderr.isatty()

    try:
        for epoch in epoch_iter:
            batches: Any = train_loader
            if show_batches:
                batches = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}", leave=False)
            train_loss, train_acc = self._run_epoch(batches, train=True)
            val_loss, val_acc = self._run_epoch(val_loader, train=False)

            if self.use_plateau_scheduler:
                self.scheduler.step(val_loss)
            else:
                self.scheduler.step()

            if use_tqdm:
                desc = f"Epoch {epoch + 1}/{epochs} | train_loss={train_loss:.4f} val_loss={val_loss:.4f}"
                if self._is_classification:
                    desc += f" val_acc={val_acc:.4f}"
                epoch_iter.set_description(desc)

            if csv_writer is not None:
                row = {
                    "epoch": epoch + 1,
                    "train_loss": f"{train_loss:.6f}",
                    "val_loss": f"{val_loss:.6f}",
                }
                if self._is_classification:
                    row["train_acc"] = f"{train_acc:.6f}"
                    row["val_acc"] = f"{val_acc:.6f}"
                csv_writer.writerow(row)
                csv_file.flush()

            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                patience_counter = 0

                if checkpoint_dir is not None:
                    ckpt_path = (
                        checkpoint_dir / f"epoch={epoch + 1:02d}-val_loss={val_loss:.4f}.pt"
                    )
                    self.save_checkpoint(str(ckpt_path), epoch=epoch + 1)
                    self.best_checkpoint_path = str(ckpt_path)
                    saved_checkpoints.append((val_loss, str(ckpt_path)))
                    saved_checkpoints.sort(key=lambda x: x[0])
                    while len(saved_checkpoints) > save_top_k:
                        _, old_path = saved_checkpoints.pop()
                        if os.path.exists(old_path):
                            os.remove(old_path)
            else:
                patience_counter += 1
                if patience_counter >= early_stopping_patience:
                    if not use_tqdm:
                        logger.info(f"Early stopping at epoch {epoch + 1}")
                    break
    finally:
        if csv_file is not None:
            csv_file.close()

    if restore_best:
        self.restore_best_checkpoint()

restore_best_checkpoint 

restore_best_checkpoint() -> bool

Reload the weights of the best-validation-loss checkpoint written during fit.

Returns:

Type Description
bool

True if a checkpoint was restored, False if none was available.
Source code in deepecgkit/training/train.py 
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
def restore_best_checkpoint(self) -> bool:
    """Reload the weights of the best-validation-loss checkpoint written during fit.

    Returns:
        True if a checkpoint was restored, False if none was available.
    """
    if self.best_checkpoint_path is None or not os.path.exists(self.best_checkpoint_path):
        if self.best_checkpoint_path is not None:
            warnings.warn(
                f"Best checkpoint {self.best_checkpoint_path} is missing; "
                "evaluation will use final-epoch weights.",
                RuntimeWarning,
                stacklevel=2,
            )
        return False

    checkpoint = load_checkpoint_file(self.best_checkpoint_path)
    self.model.load_state_dict(checkpoint["model_state_dict"])
    self.model.to(self.device)
    return True

test 

test(data_module: Any) -> dict[str, float]

Evaluate the model on the test set.

Parameters:

Name Type Description Default
data_module Any

ECGDataModule (or any object with test_dataloader)

 required

Returns:

Type Description
dict[str, float]

Dict with test_loss and test_acc (if classification)
Source code in deepecgkit/training/train.py 
386
387
388
389
390
391
392
393
394
395
396
397
398
399
def test(self, data_module: Any) -> dict[str, float]:
    """Evaluate the model on the test set.

    Args:
        data_module: ECGDataModule (or any object with test_dataloader)

    Returns:
        Dict with test_loss and test_acc (if classification)
    """
    if hasattr(data_module, "setup"):
        data_module.setup(stage="test")

    test_loader = data_module.test_dataloader()
    return self._evaluate_loader(test_loader)

validate 

validate(data_module: Any) -> dict[str, float]

Evaluate the model on the validation set.

Parameters:

Name Type Description Default
data_module Any

ECGDataModule (or any object with val_dataloader)

 required

Returns:

Type Description
dict[str, float]

Dict with val_loss and val_acc (if classification)
Source code in deepecgkit/training/train.py 
401
402
403
404
405
406
407
408
409
410
411
412
413
414
def validate(self, data_module: Any) -> dict[str, float]:
    """Evaluate the model on the validation set.

    Args:
        data_module: ECGDataModule (or any object with val_dataloader)

    Returns:
        Dict with val_loss and val_acc (if classification)
    """
    if hasattr(data_module, "setup"):
        data_module.setup(stage="validate")

    val_loader = data_module.val_dataloader()
    return self._evaluate_loader(val_loader, prefix="val")

get_test_results 

get_test_results() -> tuple[
    np.ndarray | None, np.ndarray | None, np.ndarray | None
]

Get test predictions, targets, and probabilities as numpy arrays.

Returns:

Type Description
ndarray | None

Tuple of (predictions, targets, probabilities) as numpy arrays,
ndarray | None

or (None, None, None) if no test results available.
Source code in deepecgkit/training/train.py 
470
471
472
473
474
475
476
477
478
479
480
481
482
483
def get_test_results(self) -> tuple[np.ndarray | None, np.ndarray | None, np.ndarray | None]:
    """Get test predictions, targets, and probabilities as numpy arrays.

    Returns:
        Tuple of (predictions, targets, probabilities) as numpy arrays,
        or (None, None, None) if no test results available.
    """
    if not self.test_predictions:
        return None, None, None
    return (
        torch.cat(self.test_predictions).numpy(),
        torch.cat(self.test_targets).numpy(),
        torch.cat(self.test_probabilities).numpy(),
    )

get_test_logits 

get_test_logits() -> np.ndarray | None

Get the raw model outputs from the most recent evaluation pass.

Logits are kept alongside probabilities because temperature scaling needs them and the softmax and sigmoid that produce the probabilities are not losslessly invertible in float32: a probability saturates to exactly 1.0 once its logit passes about 16.6, which is well inside the range an overconfident model reaches.

Returns:

Type Description
ndarray | None

Raw outputs as a numpy array, shaped (n_samples, output_size), or
ndarray | None

None when no evaluation has been run.
Source code in deepecgkit/training/train.py 
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
def get_test_logits(self) -> np.ndarray | None:
    """Get the raw model outputs from the most recent evaluation pass.

    Logits are kept alongside probabilities because temperature scaling
    needs them and the softmax and sigmoid that produce the probabilities
    are not losslessly invertible in float32: a probability saturates to
    exactly 1.0 once its logit passes about 16.6, which is well inside the
    range an overconfident model reaches.

    Returns:
        Raw outputs as a numpy array, shaped (n_samples, output_size), or
        None when no evaluation has been run.
    """
    if not self.test_logits:
        return None
    return torch.cat(self.test_logits).numpy()

save_checkpoint 

save_checkpoint(
    path: str | Path, epoch: int | None = None
) -> None

Save a checkpoint.

Parameters:

Name Type Description Default
path str | Path

File path to save to

 required
epoch int | None

Current epoch number (optional)

 None
Source code in deepecgkit/training/train.py 
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
def save_checkpoint(self, path: str | Path, epoch: int | None = None) -> None:
    """Save a checkpoint.

    Args:
        path: File path to save to
        epoch: Current epoch number (optional)
    """
    checkpoint: dict[str, Any] = {
        "model_state_dict": self.model.state_dict(),
        "train_config": self.train_config,
        "model_config": self.model_config,
        "data_config": self.data_config,
        "epoch": epoch,
        "best_val_loss": self.best_val_loss,
        "deepecgkit_version": _package_version(),
    }
    if self.optimizer is not None:
        checkpoint["optimizer_state_dict"] = self.optimizer.state_dict()
    torch.save(checkpoint, path)

load_checkpoint classmethod 

load_checkpoint(
    path: str | Path,
    model: Module | None = None,
    device: str | device = "auto",
) -> ECGTrainer

Load a trainer from a checkpoint.

Parameters:

Name Type Description Default
path str | Path

Path to checkpoint file

 required
model Module | None

Model instance to load weights into. When omitted, the architecture recorded in the checkpoint is rebuilt.

 None
device str | device

Device to load onto

 'auto'

Returns:

Type Description
ECGTrainer

ECGTrainer instance with loaded weights

Raises:

Type Description
ValueError

If no model is given and the checkpoint does not record which architecture it holds.
Source code in deepecgkit/training/train.py 
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
@classmethod
def load_checkpoint(
    cls,
    path: str | Path,
    model: nn.Module | None = None,
    device: str | torch.device = "auto",
) -> "ECGTrainer":
    """Load a trainer from a checkpoint.

    Args:
        path: Path to checkpoint file
        model: Model instance to load weights into. When omitted, the
            architecture recorded in the checkpoint is rebuilt.
        device: Device to load onto

    Returns:
        ECGTrainer instance with loaded weights

    Raises:
        ValueError: If no model is given and the checkpoint does not
            record which architecture it holds.
    """
    checkpoint = load_checkpoint_file(path)

    if model is None:
        from deepecgkit.models.factory import build_model_from_checkpoint

        model, _ = build_model_from_checkpoint(path)
    else:
        model.load_state_dict(checkpoint["model_state_dict"])

    trainer = cls(
        model=model,
        train_config=checkpoint.get("train_config"),
        device=device,
    )
    trainer.model_config = dict(checkpoint.get("model_config") or {})
    trainer.data_config = dict(checkpoint.get("data_config") or {})
    trainer.best_val_loss = checkpoint.get("best_val_loss", float("inf"))
    return trainer

seed_everything staticmethod 

seed_everything(seed: int) -> None

Set random seeds for reproducibility.

Parameters:

Name Type Description Default
seed int

Random seed value

 required
Source code in deepecgkit/training/train.py 
563
564
565
566
567
568
569
570
571
572
573
574
@staticmethod
def seed_everything(seed: int) -> None:
    """Set random seeds for reproducibility.

    Args:
        seed: Random seed value
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

FeatureExtractor

Bases: TransformerMixin, BaseEstimator

Runs a model's feature path over data and returns the embeddings.

The extractor is a scikit-learn transformer, so it drops straight into a Pipeline ahead of any downstream estimator. Note that sklearn's clone deep-copies the wrapped model, which doubles its memory.

Parameters:

Name Type Description Default
model Module

Any model exposing extract_features and feature_dim

 required
model_name str | None

Registry name recorded on the results, for provenance

 None
batch_size int

Samples per forward pass, ignored for a supplied DataLoader

 32
device str

Device name, or auto to prefer CUDA then MPS then CPU

 'auto'
dtype str

Precision of the returned features, float32 or float16

 'float32'
num_workers int

Worker processes for loaders this extractor constructs

 0
progress bool

Show a progress bar when tqdm is installed

 False
Example

extractor = FeatureExtractor.from_checkpoint("best.ckpt", "kanres") result = extractor.extract(signals) result.features.shape (100, 64)

Source code in deepecgkit/features.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
class FeatureExtractor(TransformerMixin, BaseEstimator):
    """Runs a model's feature path over data and returns the embeddings.

    The extractor is a scikit-learn transformer, so it drops straight into a
    Pipeline ahead of any downstream estimator. Note that sklearn's ``clone``
    deep-copies the wrapped model, which doubles its memory.

    Args:
        model: Any model exposing extract_features and feature_dim
        model_name: Registry name recorded on the results, for provenance
        batch_size: Samples per forward pass, ignored for a supplied DataLoader
        device: Device name, or auto to prefer CUDA then MPS then CPU
        dtype: Precision of the returned features, float32 or float16
        num_workers: Worker processes for loaders this extractor constructs
        progress: Show a progress bar when tqdm is installed

    Example:
        >>> extractor = FeatureExtractor.from_checkpoint("best.ckpt", "kanres")
        >>> result = extractor.extract(signals)
        >>> result.features.shape
        (100, 64)
    """

    def __init__(
        self,
        model: nn.Module,
        model_name: str | None = None,
        batch_size: int = 32,
        device: str = "auto",
        dtype: str = "float32",
        num_workers: int = 0,
        progress: bool = False,
    ) -> None:
        self.model = model
        self.model_name = model_name
        self.batch_size = batch_size
        self.device = device
        self.dtype = dtype
        self.num_workers = num_workers
        self.progress = progress

    @classmethod
    def from_checkpoint(
        cls,
        checkpoint: str | Path,
        model_name: str | None = None,
        input_channels: int | None = None,
        output_size: int | None = None,
        **kwargs: Any,
    ) -> "FeatureExtractor":
        """Build an extractor from a saved checkpoint.

        Args:
            checkpoint: Path to an ECGTrainer checkpoint, a bare state dict, or
                a legacy PyTorch Lightning checkpoint
            model_name: Registered model name matching the weights; optional
                when the checkpoint records it
            input_channels: Lead count, required for architectures whose weights
                do not encode it
            output_size: Classifier width, normally recovered from the weights
            **kwargs: Forwarded to the constructor

        Returns:
            An extractor wrapping the loaded model.
        """
        model, _ = build_model_from_checkpoint(
            checkpoint,
            name=model_name,
            input_channels=input_channels,
            output_size=output_size,
        )
        resolved = model_name or getattr(type(model), "_registry_name", None)
        return cls(model, model_name=resolved, **kwargs)

    @classmethod
    def from_pretrained(
        cls,
        weights: str | Path,
        model_name: str | None = None,
        **kwargs: Any,
    ) -> "FeatureExtractor":
        """Build an extractor from registered pretrained weights or a weights file.

        Args:
            weights: Name of registered pretrained weights, or a path
            model_name: Registered model name; inferred from the weight registry
                when the weights are registered
            **kwargs: Forwarded to the constructor, except that recognised model
                constructor arguments are passed to the model instead

        Returns:
            An extractor wrapping the loaded model.

        Raises:
            ValueError: If the model name is neither given nor inferable.
        """
        extractor_kwargs, model_kwargs = split_wrapper_kwargs(cls, kwargs)
        model, resolved = load_pretrained_resolving_name(weights, model_name, **model_kwargs)
        return cls(model, model_name=resolved, **extractor_kwargs)

    @property
    def feature_dim(self) -> int:
        """Width of the embedding the wrapped model produces.

        Raises:
            AttributeError: If the model does not declare feature_dim.
        """
        if not hasattr(self.model, "feature_dim"):
            raise AttributeError(
                f"{type(self.model).__name__} does not declare feature_dim, so it cannot be "
                f"used for feature extraction."
            )
        return int(self.model.feature_dim)

    @property
    def input_channels(self) -> int | None:
        """Lead count recovered from the weights, or None when unrecoverable."""
        return model_input_channels(self.model)

    def extract(self, data: FeatureInput, split: str = "test") -> FeatureSet:
        """Run the model's feature path over the data.

        Args:
            data: An array, tensor, file path, Dataset, DataLoader, or
                ECGDataModule
            split: Which split to use when data is an ECGDataModule, one of
                train, val, test, or all

        Returns:
            The embeddings, with labels when the input supplied them.

        Raises:
            ValueError: If the split or dtype is unknown.
            RuntimeError: If the model returns a non-2D representation.
        """
        if split not in SPLITS:
            raise ValueError(f"Unknown split: {split}. Choose one of {', '.join(SPLITS)}.")
        if self.dtype not in _DTYPES:
            raise ValueError(
                f"Unknown dtype: {self.dtype}. Choose one of {', '.join(sorted(_DTYPES))}."
            )
        if not callable(getattr(self.model, "extract_features", None)):
            raise TypeError(
                f"{type(self.model).__name__} has no extract_features() method, so its "
                f"embeddings cannot be read out. Subclass deepecgkit.models.ECGModelBase, "
                f"or add extract_features(x) returning a (batch, feature_dim) tensor."
            )

        out_dtype = _DTYPES[self.dtype]
        device = _resolve_device(self.device)
        loader = self._as_dataloader(data, split)

        was_training = self.model.training
        self.model.eval()
        self.model.to(device)

        feature_chunks: list[torch.Tensor] = []
        label_chunks: list[torch.Tensor] = []
        try:
            iterator: Any = loader
            if self.progress and tqdm is not None:
                iterator = tqdm(loader, desc="Extracting features")
            with torch.no_grad():
                for batch in iterator:
                    signals, labels = unpack_batch(batch)
                    signals = signals.to(device=device, dtype=torch.float32)
                    features = self.model.extract_features(signals)
                    if features.dim() != 2:
                        raise RuntimeError(
                            f"{type(self.model).__name__}.extract_features returned a "
                            f"{features.dim()}-D tensor; a 2-D (batch, feature_dim) tensor is "
                            f"required."
                        )
                    feature_chunks.append(features.detach().to(device="cpu", dtype=out_dtype))
                    if labels is not None:
                        label_chunks.append(labels.detach().cpu())
        finally:
            if was_training:
                self.model.train()

        if not feature_chunks:
            empty = np.empty((0, self.feature_dim), dtype=np.dtype(self.dtype))
            return FeatureSet(empty, None, self.model_name)

        features_array = torch.cat(feature_chunks, dim=0).numpy()
        labels_array = None
        if label_chunks and len(label_chunks) == len(feature_chunks):
            labels_array = torch.cat(label_chunks, dim=0).numpy()

        return FeatureSet(features_array, labels_array, self.model_name)

    def fit(self, X: Any = None, y: Any = None) -> "FeatureExtractor":  # noqa: N803
        """Return self; the extractor holds a trained model and learns nothing.

        Args:
            X: Ignored, present for scikit-learn compatibility
            y: Ignored, present for scikit-learn compatibility

        Returns:
            This extractor.
        """
        return self

    def transform(self, X: FeatureInput) -> np.ndarray:  # noqa: N803
        """Return the embeddings as a plain array, for use inside a Pipeline.

        Args:
            X: Any input accepted by extract

        Returns:
            Array of shape (num_samples, feature_dim).
        """
        return self.extract(X).features

    def get_feature_names_out(self, input_features: Any | None = None) -> np.ndarray:
        """Return names for the embedding dimensions.

        Args:
            input_features: Ignored, present for scikit-learn compatibility

        Returns:
            Array of feature names.
        """
        prefix = self.model_name or type(self.model).__name__.lower()
        return np.asarray([f"{prefix}_{i}" for i in range(self.feature_dim)], dtype=object)

    def _as_dataloader(self, data: FeatureInput, split: str) -> DataLoader:
        """Normalise any supported input to a non-shuffling DataLoader."""
        return as_inference_dataloader(
            data,
            split=split,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            input_channels=self.input_channels,
            verb="extracted",
        )

feature_dim property 

feature_dim: int

Width of the embedding the wrapped model produces.

Raises:

Type Description
AttributeError

If the model does not declare feature_dim.

input_channels property 

input_channels: int | None

Lead count recovered from the weights, or None when unrecoverable.

from_checkpoint classmethod 

from_checkpoint(
    checkpoint: str | Path,
    model_name: str | None = None,
    input_channels: int | None = None,
    output_size: int | None = None,
    **kwargs: Any,
) -> FeatureExtractor

Build an extractor from a saved checkpoint.

Parameters:

Name Type Description Default
checkpoint str | Path

Path to an ECGTrainer checkpoint, a bare state dict, or a legacy PyTorch Lightning checkpoint

 required
model_name str | None

Registered model name matching the weights; optional when the checkpoint records it

 None
input_channels int | None

Lead count, required for architectures whose weights do not encode it

 None
output_size int | None

Classifier width, normally recovered from the weights

 None
**kwargs Any

Forwarded to the constructor

 {}

Returns:

Type Description
FeatureExtractor

An extractor wrapping the loaded model.
Source code in deepecgkit/features.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
@classmethod
def from_checkpoint(
    cls,
    checkpoint: str | Path,
    model_name: str | None = None,
    input_channels: int | None = None,
    output_size: int | None = None,
    **kwargs: Any,
) -> "FeatureExtractor":
    """Build an extractor from a saved checkpoint.

    Args:
        checkpoint: Path to an ECGTrainer checkpoint, a bare state dict, or
            a legacy PyTorch Lightning checkpoint
        model_name: Registered model name matching the weights; optional
            when the checkpoint records it
        input_channels: Lead count, required for architectures whose weights
            do not encode it
        output_size: Classifier width, normally recovered from the weights
        **kwargs: Forwarded to the constructor

    Returns:
        An extractor wrapping the loaded model.
    """
    model, _ = build_model_from_checkpoint(
        checkpoint,
        name=model_name,
        input_channels=input_channels,
        output_size=output_size,
    )
    resolved = model_name or getattr(type(model), "_registry_name", None)
    return cls(model, model_name=resolved, **kwargs)

from_pretrained classmethod 

from_pretrained(
    weights: str | Path,
    model_name: str | None = None,
    **kwargs: Any,
) -> FeatureExtractor

Build an extractor from registered pretrained weights or a weights file.

Parameters:

Name Type Description Default
weights str | Path

Name of registered pretrained weights, or a path

 required
model_name str | None

Registered model name; inferred from the weight registry when the weights are registered

 None
**kwargs Any

Forwarded to the constructor, except that recognised model constructor arguments are passed to the model instead

 {}

Returns:

Type Description
FeatureExtractor

An extractor wrapping the loaded model.

Raises:

Type Description
ValueError

If the model name is neither given nor inferable.
Source code in deepecgkit/features.py 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
@classmethod
def from_pretrained(
    cls,
    weights: str | Path,
    model_name: str | None = None,
    **kwargs: Any,
) -> "FeatureExtractor":
    """Build an extractor from registered pretrained weights or a weights file.

    Args:
        weights: Name of registered pretrained weights, or a path
        model_name: Registered model name; inferred from the weight registry
            when the weights are registered
        **kwargs: Forwarded to the constructor, except that recognised model
            constructor arguments are passed to the model instead

    Returns:
        An extractor wrapping the loaded model.

    Raises:
        ValueError: If the model name is neither given nor inferable.
    """
    extractor_kwargs, model_kwargs = split_wrapper_kwargs(cls, kwargs)
    model, resolved = load_pretrained_resolving_name(weights, model_name, **model_kwargs)
    return cls(model, model_name=resolved, **extractor_kwargs)

extract 

extract(
    data: FeatureInput, split: str = "test"
) -> FeatureSet

Run the model's feature path over the data.

Parameters:

Name Type Description Default
data FeatureInput

An array, tensor, file path, Dataset, DataLoader, or ECGDataModule

 required
split str

Which split to use when data is an ECGDataModule, one of train, val, test, or all

 'test'

Returns:

Type Description
FeatureSet

The embeddings, with labels when the input supplied them.

Raises:

Type Description
ValueError

If the split or dtype is unknown.
RuntimeError

If the model returns a non-2D representation.
Source code in deepecgkit/features.py 
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
def extract(self, data: FeatureInput, split: str = "test") -> FeatureSet:
    """Run the model's feature path over the data.

    Args:
        data: An array, tensor, file path, Dataset, DataLoader, or
            ECGDataModule
        split: Which split to use when data is an ECGDataModule, one of
            train, val, test, or all

    Returns:
        The embeddings, with labels when the input supplied them.

    Raises:
        ValueError: If the split or dtype is unknown.
        RuntimeError: If the model returns a non-2D representation.
    """
    if split not in SPLITS:
        raise ValueError(f"Unknown split: {split}. Choose one of {', '.join(SPLITS)}.")
    if self.dtype not in _DTYPES:
        raise ValueError(
            f"Unknown dtype: {self.dtype}. Choose one of {', '.join(sorted(_DTYPES))}."
        )
    if not callable(getattr(self.model, "extract_features", None)):
        raise TypeError(
            f"{type(self.model).__name__} has no extract_features() method, so its "
            f"embeddings cannot be read out. Subclass deepecgkit.models.ECGModelBase, "
            f"or add extract_features(x) returning a (batch, feature_dim) tensor."
        )

    out_dtype = _DTYPES[self.dtype]
    device = _resolve_device(self.device)
    loader = self._as_dataloader(data, split)

    was_training = self.model.training
    self.model.eval()
    self.model.to(device)

    feature_chunks: list[torch.Tensor] = []
    label_chunks: list[torch.Tensor] = []
    try:
        iterator: Any = loader
        if self.progress and tqdm is not None:
            iterator = tqdm(loader, desc="Extracting features")
        with torch.no_grad():
            for batch in iterator:
                signals, labels = unpack_batch(batch)
                signals = signals.to(device=device, dtype=torch.float32)
                features = self.model.extract_features(signals)
                if features.dim() != 2:
                    raise RuntimeError(
                        f"{type(self.model).__name__}.extract_features returned a "
                        f"{features.dim()}-D tensor; a 2-D (batch, feature_dim) tensor is "
                        f"required."
                    )
                feature_chunks.append(features.detach().to(device="cpu", dtype=out_dtype))
                if labels is not None:
                    label_chunks.append(labels.detach().cpu())
    finally:
        if was_training:
            self.model.train()

    if not feature_chunks:
        empty = np.empty((0, self.feature_dim), dtype=np.dtype(self.dtype))
        return FeatureSet(empty, None, self.model_name)

    features_array = torch.cat(feature_chunks, dim=0).numpy()
    labels_array = None
    if label_chunks and len(label_chunks) == len(feature_chunks):
        labels_array = torch.cat(label_chunks, dim=0).numpy()

    return FeatureSet(features_array, labels_array, self.model_name)

fit 

fit(X: Any = None, y: Any = None) -> FeatureExtractor

Return self; the extractor holds a trained model and learns nothing.

Parameters:

Name Type Description Default
X Any

Ignored, present for scikit-learn compatibility

 None
y Any

Ignored, present for scikit-learn compatibility

 None

Returns:

Type Description
FeatureExtractor

This extractor.
Source code in deepecgkit/features.py 
320
321
322
323
324
325
326
327
328
329
330
def fit(self, X: Any = None, y: Any = None) -> "FeatureExtractor":  # noqa: N803
    """Return self; the extractor holds a trained model and learns nothing.

    Args:
        X: Ignored, present for scikit-learn compatibility
        y: Ignored, present for scikit-learn compatibility

    Returns:
        This extractor.
    """
    return self

transform 

transform(X: FeatureInput) -> np.ndarray

Return the embeddings as a plain array, for use inside a Pipeline.

Parameters:

Name Type Description Default
X FeatureInput

Any input accepted by extract

 required

Returns:

Type Description
ndarray

Array of shape (num_samples, feature_dim).
Source code in deepecgkit/features.py 
332
333
334
335
336
337
338
339
340
341
def transform(self, X: FeatureInput) -> np.ndarray:  # noqa: N803
    """Return the embeddings as a plain array, for use inside a Pipeline.

    Args:
        X: Any input accepted by extract

    Returns:
        Array of shape (num_samples, feature_dim).
    """
    return self.extract(X).features

get_feature_names_out 

get_feature_names_out(
    input_features: Any | None = None,
) -> np.ndarray

Return names for the embedding dimensions.

Parameters:

Name Type Description Default
input_features Any | None

Ignored, present for scikit-learn compatibility

 None

Returns:

Type Description
ndarray

Array of feature names.
Source code in deepecgkit/features.py 
343
344
345
346
347
348
349
350
351
352
353
def get_feature_names_out(self, input_features: Any | None = None) -> np.ndarray:
    """Return names for the embedding dimensions.

    Args:
        input_features: Ignored, present for scikit-learn compatibility

    Returns:
        Array of feature names.
    """
    prefix = self.model_name or type(self.model).__name__.lower()
    return np.asarray([f"{prefix}_{i}" for i in range(self.feature_dim)], dtype=object)

FeatureSet dataclass

Embeddings extracted from a model, with any aligned labels.

Attributes:

Name Type Description
features ndarray

Array of shape (num_samples, feature_dim)
labels ndarray | None

Aligned labels of shape (num_samples,) for single-label data or (num_samples, num_classes) for multi-label data, or None when the input carried no labels
model_name str | None

Registry name of the model that produced the features
Source code in deepecgkit/features.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
@dataclass
class FeatureSet:
    """Embeddings extracted from a model, with any aligned labels.

    Attributes:
        features: Array of shape (num_samples, feature_dim)
        labels: Aligned labels of shape (num_samples,) for single-label data or
            (num_samples, num_classes) for multi-label data, or None when the
            input carried no labels
        model_name: Registry name of the model that produced the features
    """

    features: np.ndarray
    labels: np.ndarray | None = None
    model_name: str | None = field(default=None)

    @property
    def feature_dim(self) -> int:
        """Width of the embedding."""
        return int(self.features.shape[1])

    def __len__(self) -> int:
        """Number of samples."""
        return int(self.features.shape[0])

    def __array__(self, dtype: Any | None = None) -> np.ndarray:
        """Expose the features to numpy so the object can stand in for an array."""
        if dtype is None:
            return self.features
        return self.features.astype(dtype)

    def save(self, path: str | Path) -> None:
        """Write the features to disk, choosing the format from the suffix.

        Args:
            path: Destination ending in .npz, .npy, or .csv. The .npz form
                keeps the labels and metadata; .npy stores features only; .csv
                writes one column per feature plus a label column.

        Raises:
            ValueError: If the suffix is not one of the three supported forms.
        """
        out_path = Path(path)
        out_path.parent.mkdir(parents=True, exist_ok=True)

        if out_path.suffix == ".npy":
            np.save(str(out_path), self.features)
            return

        if out_path.suffix == ".npz":
            payload: dict[str, Any] = {
                "features": self.features,
                "feature_dim": np.asarray(self.feature_dim),
            }
            if self.labels is not None:
                payload["labels"] = self.labels
            if self.model_name is not None:
                payload["model_name"] = np.asarray(self.model_name)
            np.savez(str(out_path), **payload)
            return

        if out_path.suffix == ".csv":
            import pandas as pd

            frame = pd.DataFrame(self.features, columns=[f"f{i}" for i in range(self.feature_dim)])
            if self.labels is not None:
                labels = self.labels
                if labels.ndim == 1:
                    frame["label"] = labels
                else:
                    for i in range(labels.shape[1]):
                        frame[f"label_{i}"] = labels[:, i]
            frame.to_csv(str(out_path), index=False)
            return

        raise ValueError(
            f"Unsupported output format: {out_path.suffix}. Supported formats: .npz, .npy, .csv"
        )

feature_dim property 

feature_dim: int

Width of the embedding.

save 

save(path: str | Path) -> None

Write the features to disk, choosing the format from the suffix.

Parameters:

Name Type Description Default
path str | Path

Destination ending in .npz, .npy, or .csv. The .npz form keeps the labels and metadata; .npy stores features only; .csv writes one column per feature plus a label column.

 required

Raises:

Type Description
ValueError

If the suffix is not one of the three supported forms.
Source code in deepecgkit/features.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def save(self, path: str | Path) -> None:
    """Write the features to disk, choosing the format from the suffix.

    Args:
        path: Destination ending in .npz, .npy, or .csv. The .npz form
            keeps the labels and metadata; .npy stores features only; .csv
            writes one column per feature plus a label column.

    Raises:
        ValueError: If the suffix is not one of the three supported forms.
    """
    out_path = Path(path)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    if out_path.suffix == ".npy":
        np.save(str(out_path), self.features)
        return

    if out_path.suffix == ".npz":
        payload: dict[str, Any] = {
            "features": self.features,
            "feature_dim": np.asarray(self.feature_dim),
        }
        if self.labels is not None:
            payload["labels"] = self.labels
        if self.model_name is not None:
            payload["model_name"] = np.asarray(self.model_name)
        np.savez(str(out_path), **payload)
        return

    if out_path.suffix == ".csv":
        import pandas as pd

        frame = pd.DataFrame(self.features, columns=[f"f{i}" for i in range(self.feature_dim)])
        if self.labels is not None:
            labels = self.labels
            if labels.ndim == 1:
                frame["label"] = labels
            else:
                for i in range(labels.shape[1]):
                    frame[f"label_{i}"] = labels[:, i]
        frame.to_csv(str(out_path), index=False)
        return

    raise ValueError(
        f"Unsupported output format: {out_path.suffix}. Supported formats: .npz, .npy, .csv"
    )

KanResWideX

Bases: ECGModelBase

KanRes-Wide-X model for ECG signal classification.

A convolutional neural network architecture designed for ECG signal analysis with residual connections and wide blocks for improved feature extraction.

Parameters:

Name Type Description Default
input_channels int

Number of input channels (default: 1 for single-lead ECG)

 1
output_size int

Number of output classes or regression targets

 4
base_channels int

Base number of channels for the first layer (default: 64)

 64
Example

model = KanResWideX(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape) # [32, 4]

features = model.extract_features(x) print(features.shape) # (32, 64)

Source code in deepecgkit/models/kanres_wide_x.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
@register_model(
    name="kanres",
    description="KAN-ResNet architecture with wide layers",
)
class KanResWideX(ECGModelBase):
    """
    KanRes-Wide-X model for ECG signal classification.

    A convolutional neural network architecture designed for ECG signal analysis
    with residual connections and wide blocks for improved feature extraction.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes or regression targets
        base_channels: Base number of channels for the first layer (default: 64)

    Example:
        >>> model = KanResWideX(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)  # [32, 4]

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 64)
    """

    def __init__(self, input_channels: int = 1, output_size: int = 4, base_channels: int = 64):
        super().__init__()

        self.input_layer = ConvBlock(input_channels, base_channels)
        self.res_modules = nn.Sequential(
            KanResModule(base_channels), KanResModule(base_channels), KanResModule(base_channels)
        )
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self._feature_dim = base_channels
        self.classifier = nn.Linear(base_channels, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_layer(x)
        x = self.res_modules(x)
        x = self.global_pool(x)
        x = x.squeeze(-1)
        return x

extract_features 

extract_features(
    data: FeatureInput,
    model: Module | None = None,
    checkpoint: str | Path | None = None,
    weights: str | Path | None = None,
    model_name: str | None = None,
    input_channels: int | None = None,
    output_size: int | None = None,
    split: str = "test",
    **kwargs: Any,
) -> FeatureSet

Extract deep embeddings in one call.

Exactly one weight source must be given: an already-built model, a checkpoint path, or pretrained weights.

Parameters:

Name Type Description Default
data FeatureInput

An array, tensor, file path, Dataset, DataLoader, or ECGDataModule

 required
model Module | None

An already-built model

 None
checkpoint str | Path | None

Path to a checkpoint

 None
weights str | Path | None

Registered pretrained weight name, or a path to a weights file

 None
model_name str | None

Registered model name; optional for a checkpoint that records its own

 None
input_channels int | None

Lead count, when it cannot be recovered from the weights

 None
output_size int | None

Classifier width, when it cannot be recovered

 None
split str

Which split to use when data is an ECGDataModule

 'test'
**kwargs Any

Forwarded to FeatureExtractor

 {}

Returns:

Type Description
FeatureSet

The embeddings, with labels when the input supplied them.

Raises:

Type Description
ValueError

If not exactly one weight source is given, or a checkpoint neither records its architecture nor is given a model name.
Example

result = extract_features(signals, checkpoint="best.ckpt", model_name="kanres") result.features.shape (100, 64)

Source code in deepecgkit/features.py 
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
def extract_features(
    data: FeatureInput,
    model: nn.Module | None = None,
    checkpoint: str | Path | None = None,
    weights: str | Path | None = None,
    model_name: str | None = None,
    input_channels: int | None = None,
    output_size: int | None = None,
    split: str = "test",
    **kwargs: Any,
) -> FeatureSet:
    """Extract deep embeddings in one call.

    Exactly one weight source must be given: an already-built model, a
    checkpoint path, or pretrained weights.

    Args:
        data: An array, tensor, file path, Dataset, DataLoader, or ECGDataModule
        model: An already-built model
        checkpoint: Path to a checkpoint
        weights: Registered pretrained weight name, or a path to a weights file
        model_name: Registered model name; optional for a checkpoint that
            records its own
        input_channels: Lead count, when it cannot be recovered from the weights
        output_size: Classifier width, when it cannot be recovered
        split: Which split to use when data is an ECGDataModule
        **kwargs: Forwarded to FeatureExtractor

    Returns:
        The embeddings, with labels when the input supplied them.

    Raises:
        ValueError: If not exactly one weight source is given, or a checkpoint
            neither records its architecture nor is given a model name.

    Example:
        >>> result = extract_features(signals, checkpoint="best.ckpt", model_name="kanres")
        >>> result.features.shape
        (100, 64)
    """
    supplied = dict(_iter_inputs(model=model, checkpoint=checkpoint, weights=weights))
    if len(supplied) != 1:
        raise ValueError(
            f"Pass exactly one of model, checkpoint, or weights; got "
            f"{', '.join(sorted(supplied)) or 'none'}."
        )

    if model is not None:
        extractor = FeatureExtractor(model, model_name=model_name, **kwargs)
    elif checkpoint is not None:
        extractor = FeatureExtractor.from_checkpoint(
            checkpoint,
            model_name,
            input_channels=input_channels,
            output_size=output_size,
            **kwargs,
        )
    else:
        extractor = FeatureExtractor.from_pretrained(weights, model_name=model_name, **kwargs)

    return extractor.extract(data, split=split)

read_csv 

read_csv(
    csv_file: str,
    delimiter: str = ",",
    transpose: bool = False,
    skip_header: bool = True,
    dtype: type | None = None,
) -> tuple[np.ndarray, dict[str, int]]

Read CSV file and return data array and header mapping.

Parameters:

Name Type Description Default
csv_file str

Path to the CSV file

 required
delimiter str

Column delimiter (default: ',')

 ','
transpose bool

Whether to transpose the data array

 False
skip_header bool

Whether to skip the first row as header

 True
dtype type | None

Data type for the numpy array

 None

Returns:

Type Description
tuple[ndarray, dict[str, int]]

Tuple of (data_array, header_mapping)
Source code in deepecgkit/utils/__init__.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def read_csv(
    csv_file: str,
    delimiter: str = ",",
    transpose: bool = False,
    skip_header: bool = True,
    dtype: type | None = None,
) -> tuple[np.ndarray, dict[str, int]]:
    """
    Read CSV file and return data array and header mapping.

    Args:
        csv_file: Path to the CSV file
        delimiter: Column delimiter (default: ',')
        transpose: Whether to transpose the data array
        skip_header: Whether to skip the first row as header
        dtype: Data type for the numpy array

    Returns:
        Tuple of (data_array, header_mapping)
    """
    if not os.path.exists(csv_file):
        raise FileNotFoundError(f"CSV file not found: {csv_file}")

    data: list[list[str]] = []
    header: dict[str, int] = {}

    try:
        with open(csv_file) as f:
            csv_data = csv.reader(f, delimiter=delimiter)

            if skip_header:
                try:
                    temp = next(csv_data)
                    header = {k: v for v, k in enumerate(temp)}
                except StopIteration as err:
                    raise ValueError("File is empty") from err

            for row in csv_data:
                data.append(row)

        if not data:
            raise ValueError("No data found in CSV file")

        data_array = np.array(data, dtype=dtype)

        if transpose:
            data_array = np.transpose(data_array)

        return data_array, header

    except Exception as err:
        raise ValueError(f"Error reading CSV file: {err!s}") from err