Skip to content

datasets

Classes

Functions

create_data_pipeline 

create_data_pipeline(
    ds: Dataset,
    sampling_rate: int,
    batch_size: int,
    buffer_size: int | None = None,
    augmentations: list[NamedParams] | None = None,
    num_classes: int = 2,
) -> tf.data.Dataset

Create data pipeline for training

Parameters:

  • ds

    (Dataset) –

    Dataset

  • sampling_rate

    (int) –

    Sampling rate

  • batch_size

    (int) –

    Batch size

  • buffer_size

    (int, default: None ) –

    Buffer size. Defaults to None.

  • augmentations

    (list[NamedParams], default: None ) –

    Augmentations. Defaults to None.

  • num_classes

    (int, default: 2 ) –

    Number of classes. Defaults to 2.

Returns:

  • Dataset –

    tf.data.Dataset: Data pipeline

Source code in heartkit/tasks/rhythm/datasets.py 
15
16
17
18
19
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
def create_data_pipeline(
    ds: tf.data.Dataset,
    sampling_rate: int,
    batch_size: int,
    buffer_size: int | None = None,
    augmentations: list[NamedParams] | None = None,
    num_classes: int = 2,
) -> tf.data.Dataset:
    """Create data pipeline for training

    Args:
        ds (tf.data.Dataset): Dataset
        sampling_rate (int): Sampling rate
        batch_size (int): Batch size
        buffer_size (int, optional): Buffer size. Defaults to None.
        augmentations (list[NamedParams], optional): Augmentations. Defaults to None.
        num_classes (int, optional): Number of classes. Defaults to 2.

    Returns:
        tf.data.Dataset: Data pipeline
    """
    if buffer_size:
        ds = ds.shuffle(
            buffer_size=buffer_size,
            reshuffle_each_iteration=True,
        )
    if batch_size:
        ds = ds.batch(
            batch_size=batch_size,
            drop_remainder=True,
            num_parallel_calls=tf.data.AUTOTUNE,
        )
    augmenter = create_augmentation_pipeline(augmentations, sampling_rate=sampling_rate)

    ds = (
        ds.map(
            lambda data, labels: {
                "data": tf.cast(data, "float32"),
                "labels": tf.one_hot(labels, num_classes),
            },
            num_parallel_calls=tf.data.AUTOTUNE,
        )
        .map(
            augmenter,
            num_parallel_calls=tf.data.AUTOTUNE,
        )
        .map(
            lambda data: (data["data"], data["labels"]),
            num_parallel_calls=tf.data.AUTOTUNE,
        )
    )
    return ds.prefetch(tf.data.AUTOTUNE)

load_train_datasets 

load_train_datasets(datasets: list[HKDataset], params: HKTaskParams) -> tuple[tf.data.Dataset, tf.data.Dataset]

Load training and validation datasets

Parameters:

Returns:

  • tuple[Dataset, Dataset] –

    tuple[tf.data.Dataset, tf.data.Dataset]: Training and validation datasets

Source code in heartkit/tasks/rhythm/datasets.py 
 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
def load_train_datasets(
    datasets: list[HKDataset],
    params: HKTaskParams,
) -> tuple[tf.data.Dataset, tf.data.Dataset]:
    """Load training and validation datasets

    Args:
        datasets (list[HKDataset]): List of datasets
        params (HKTaskParams): Task parameters

    Returns:
        tuple[tf.data.Dataset, tf.data.Dataset]: Training and validation datasets
    """

    train_ds, val_ds = load_train_dataloader_split(datasets, params, factory=RhythmDataloaderFactory)

    # Shuffle and batch datasets for training
    train_ds = create_data_pipeline(
        ds=train_ds,
        sampling_rate=params.sampling_rate,
        batch_size=params.batch_size,
        buffer_size=params.buffer_size,
        augmentations=params.augmentations + params.preprocesses,
        num_classes=params.num_classes,
    )

    val_ds = create_data_pipeline(
        ds=val_ds,
        sampling_rate=params.sampling_rate,
        batch_size=params.batch_size,
        augmentations=params.preprocesses,
        num_classes=params.num_classes,
    )

    # Cache validation dataset w/ fixed size
    val_steps_per_epoch = params.val_size // params.batch_size if params.val_size else params.val_steps_per_epoch
    logger.info(f"Validation steps per epoch: {val_steps_per_epoch}")
    val_ds = val_ds.take(val_steps_per_epoch).cache()

    return train_ds, val_ds

load_test_dataset 

load_test_dataset(datasets: list[HKDataset], params: HKTaskParams) -> tf.data.Dataset

Load test dataset

Parameters:

Returns:

  • Dataset –

    tf.data.Dataset: Test dataset

Source code in heartkit/tasks/rhythm/datasets.py 
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
def load_test_dataset(
    datasets: list[HKDataset],
    params: HKTaskParams,
) -> tf.data.Dataset:
    """Load test dataset

    Args:
        datasets (list[HKDataset]): List of datasets
        params (HKTaskParams): Task parameters

    Returns:
        tf.data.Dataset: Test dataset
    """

    test_ds = load_test_dataloader_split(datasets, params, factory=RhythmDataloaderFactory)

    test_ds = create_data_pipeline(
        ds=test_ds,
        sampling_rate=params.sampling_rate,
        batch_size=params.batch_size,
        augmentations=params.preprocesses,
        num_classes=params.num_classes,
    )

    test_ds = test_ds.take(params.test_size // params.batch_size).cache()

    return test_ds