Stats Loader#

loader #

FUNCTION DESCRIPTION
loader

Create dataloader for computing running statistics with automatic GPU cleanup.

Classes#

FixedSubsetSampler(samples: list[int]) #

Bases: Sampler

Sampler that yields a fixed sequence of dataset indices.

Provides deterministic iteration over a predefined subset of indices. Supports creating nested subsets via index translation.

PARAMETER DESCRIPTION
samples

Fixed list of dataset indices to sample.

TYPE: list[int]

ATTRIBUTE DESCRIPTION
samples

The fixed sequence of indices.

TYPE: list[int]

Examples:

>>> from spectre.stats.loader import FixedSubsetSampler
>>> from torch.utils.data import DataLoader, TensorDataset
>>> import torch
>>>
>>> X = torch.randn(100, 10)
>>> dataset = TensorDataset(X)
>>>
>>> # Sample indices [5, 10, 15, 20]
>>> sampler = FixedSubsetSampler([5, 10, 15, 20])
>>> loader = DataLoader(dataset, sampler=sampler)
>>> len(loader)
4
METHOD DESCRIPTION
subset

Create a sub-sampler from this sampler's indices.
dereference

Translate subset indices to original dataset indices.
Source code in spectre/stats/loader.py 
282
283
def __init__(self, samples: list[int]):
    self.samples = samples
Functions#
subset(new_subset: list[int]) #

Create a sub-sampler from this sampler's indices.

PARAMETER DESCRIPTION
new_subset

Indices into the current sampler's sequence.

TYPE: list[int]

RETURNS DESCRIPTION
FixedSubsetSampler

New sampler with dereferenced indices.
Source code in spectre/stats/loader.py 
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
def subset(self, new_subset: list[int]):
    """
    Create a sub-sampler from this sampler's indices.

    Parameters
    ----------
    new_subset : list[int]
        Indices into the current sampler's sequence.

    Returns
    -------
    FixedSubsetSampler
        New sampler with dereferenced indices.
    """
    return FixedSubsetSampler(self.dereference(new_subset))
dereference(indices: list[int]) -> list[int] #

Translate subset indices to original dataset indices.

Maps small indices (positions in current subset) to large indices (positions in original full dataset).

PARAMETER DESCRIPTION
indices

Indices into the current subset.

TYPE: list[int]

RETURNS DESCRIPTION
list[int]

Corresponding indices in the original dataset.

Examples:

>>> sampler = FixedSubsetSampler([10, 20, 30, 40])
>>> sampler.dereference([0, 2])  # Get 1st and 3rd elements
[10, 30]
Source code in spectre/stats/loader.py 
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
def dereference(self, indices: list[int]) -> list[int]:
    """
    Translate subset indices to original dataset indices.

    Maps small indices (positions in current subset) to large indices
    (positions in original full dataset).

    Parameters
    ----------
    indices : list[int]
        Indices into the current subset.

    Returns
    -------
    list[int]
        Corresponding indices in the original dataset.

    Examples
    --------
    >>> sampler = FixedSubsetSampler([10, 20, 30, 40])
    >>> sampler.dereference([0, 2])  # Get 1st and 3rd elements
    [10, 30]
    """
    return [self.samples[i] for i in indices]

FixedRandomSubsetSampler(data: Dataset, start: int | None = None, end: int | None = None, seed: int = 1) #

Bases: FixedSubsetSampler

Deterministic random sampler for dataset subsets.

Creates a fixed random permutation of dataset indices using a specified seed, then samples a contiguous slice. Provides reproducible random sampling across runs.

PARAMETER DESCRIPTION
data

The dataset to sample from (used only for length).

TYPE: Dataset

start

Starting index in the shuffled sequence, by default None (start from 0).

TYPE: int | None DEFAULT: None

end

Ending index in the shuffled sequence, by default None (go to end).

TYPE: int | None DEFAULT: None

seed

Random seed for reproducible shuffling, by default 1.

TYPE: int DEFAULT: 1

Examples:

>>> from spectre.stats.loader import FixedRandomSubsetSampler
>>> from torch.utils.data import DataLoader, TensorDataset
>>> import torch
>>>
>>> X = torch.randn(1000, 10)
>>> dataset = TensorDataset(X)
>>>
>>> # Sample 100 random indices deterministically
>>> sampler = FixedRandomSubsetSampler(dataset, end=100, seed=42)
>>> loader = DataLoader(dataset, sampler=sampler)
>>> len(loader)
100
>>>
>>> # Same seed gives same samples
>>> sampler2 = FixedRandomSubsetSampler(dataset, end=100, seed=42)
>>> list(sampler) == list(sampler2)
True
METHOD DESCRIPTION
class_subset

Filter samples by class label or custom rule.
Source code in spectre/stats/loader.py 
379
380
381
382
383
384
385
386
387
388
389
390
391
def __init__(
    self,
    data: Dataset,
    start: int | None = None,
    end: int | None = None,
    seed: int = 1,
):
    rng = random.Random(seed)
    shuffled = list(range(len(data)))
    rng.shuffle(shuffled)
    self.data = data

    super(FixedRandomSubsetSampler, self).__init__(shuffled[start:end])
Functions#
class_subset(class_filter: int | Callable) #

Filter samples by class label or custom rule.

Creates a new sampler containing only samples that match the specified class or satisfy a custom filtering function.

PARAMETER DESCRIPTION
class_filter

If int, filters for samples with that class label (assumes dataset returns (data, label) tuples). If callable, applies as filter function to each dataset item.

TYPE: int | Callable

RETURNS DESCRIPTION
FixedSubsetSampler

New sampler with filtered indices.

Examples:

>>> # Assuming dataset returns (data, label) tuples
>>> sampler = FixedRandomSubsetSampler(dataset, end=1000, seed=42)
>>>
>>> # Get only samples from class 0
>>> class0_sampler = sampler.class_subset(0)
>>>
>>> # Custom filter function
>>> def is_positive(item):
...     data, label = item
...     return label > 0
>>> positive_sampler = sampler.class_subset(is_positive)
Source code in spectre/stats/loader.py 
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
def class_subset(self, class_filter: int | Callable):
    """
    Filter samples by class label or custom rule.

    Creates a new sampler containing only samples that match the specified
    class or satisfy a custom filtering function.


    Parameters
    ----------
    class_filter : int | Callable
        If int, filters for samples with that class label (assumes dataset
        returns (data, label) tuples). If callable, applies as filter function
        to each dataset item.

    Returns
    -------
    FixedSubsetSampler
        New sampler with filtered indices.

    Examples
    --------
    >>> # Assuming dataset returns (data, label) tuples
    >>> sampler = FixedRandomSubsetSampler(dataset, end=1000, seed=42)
    >>>
    >>> # Get only samples from class 0
    >>> class0_sampler = sampler.class_subset(0)
    >>>
    >>> # Custom filter function
    >>> def is_positive(item):
    ...     data, label = item
    ...     return label > 0
    >>> positive_sampler = sampler.class_subset(is_positive)
    """
    if isinstance(class_filter, int):

        def rule(d):
            return d[1] == class_filter

    else:
        rule = class_filter
    return self.subset(
        [i for i, j in enumerate(self.samples) if rule(self.data_source[j])]
    )

Functions#

loader(stat, dataset: Union[torch.Tensor, DatasetWeighted], **kwargs) -> DataLoader #

Create dataloader for computing running statistics with automatic GPU cleanup.

Wraps a dataset with a DataLoader and automatically moves the statistic back to CPU after all batches are processed. Supports subsampling for large datasets. Automatically handles Batch objects from DatasetWeighted for weighted statistics.

PARAMETER DESCRIPTION
stat

The statistic object to compute (RunningMean, RunningVariance, etc.). Automatically moved to CPU after processing.

TYPE: RunningStats

dataset

Data tensor or DatasetWeighted to compute statistics over. First dimension is batch/sample dimension.

TYPE: Tensor | DatasetWeighted

**kwargs

Additional arguments passed to DataLoader:

  • batch_size : int - Batch size for processing (default 1)
  • num_workers : int - Number of worker processes
  • pin_memory : bool - Pin memory for faster GPU transfer
  • sample_size : int - Number of samples to process (subsampling)
  • random_seed : int - Seed for random subsampling

DEFAULT: {}

RETURNS DESCRIPTION
iterator

Iterator yielding batches or Batch objects. After exhaustion, stat is moved to CPU.
RAISES DESCRIPTION
TypeError

If stat is not a RunningStats object.

Examples:

Basic usage with tensor:

>>> from spectre.stats import RunningMean, loader
>>> import torch
>>>
>>> X = torch.randn(10000, 50)
>>> stat = RunningMean()
>>>
>>> for batch in loader(stat, X, batch_size=100):
...     stat.add(batch)
>>>
>>> mean = stat.mean()  # stat is now on CPU

Using with DatasetWeighted for weighted statistics:

>>> from spectre.data import DatasetWeighted
>>> X = torch.randn(10000, 50)
>>> weights = torch.rand(10000)
>>> dataset = DatasetWeighted(X, weights=weights)
>>> stat = RunningMean()
>>>
>>> for batch in loader(stat, dataset, batch_size=100):
...     stat.add(batch.data, weights=batch.weights)
>>>
>>> mean = stat.mean()

Subsampling for faster computation:

>>> # Process only 1000 samples instead of full dataset
>>> stat = RunningMean()
>>> for batch in loader(stat, X, batch_size=100, sample_size=1000):
...     stat.add(batch)
>>> mean = stat.mean()

Computing multiple statistics together:

>>> from spectre.stats import RunningStat, RunningMean, RunningVariance
>>>
>>> stats = RunningStat(mean=RunningMean(), var=RunningVariance())
>>> for batch in loader(stats, X, batch_size=100):
...     stats.add(batch)
>>>
>>> mean = stats.mean.mean()
>>> variance = stats.var.variance()
Notes

The loader automatically moves stat to CPU after processing to avoid memory leaks when working with GPU tensors.

When using DatasetWeighted, the loader yields Batch objects with data, weights, and optional target. Use batch.data and batch.weights when calling stat.add().

Source code in spectre/stats/loader.py 
 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
 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
def loader(stat, dataset: Union[torch.Tensor, DatasetWeighted], **kwargs) -> DataLoader:
    """
    Create dataloader for computing running statistics with automatic GPU cleanup.

    Wraps a dataset with a DataLoader and automatically moves the statistic back
    to CPU after all batches are processed. Supports subsampling for large datasets.
    Automatically handles Batch objects from DatasetWeighted for weighted statistics.

    Parameters
    ----------
    stat : RunningStats
        The statistic object to compute (RunningMean, RunningVariance, etc.).
        Automatically moved to CPU after processing.

    dataset : torch.Tensor | DatasetWeighted
        Data tensor or DatasetWeighted to compute statistics over. First dimension
        is batch/sample dimension.

    **kwargs
        Additional arguments passed to DataLoader:

        - batch_size : int - Batch size for processing (default 1)
        - num_workers : int - Number of worker processes
        - pin_memory : bool - Pin memory for faster GPU transfer
        - sample_size : int - Number of samples to process (subsampling)
        - random_seed : int - Seed for random subsampling

    Returns
    -------
    iterator
        Iterator yielding batches or Batch objects. After exhaustion, stat
        is moved to CPU.

    Raises
    ------
    TypeError
        If stat is not a RunningStats object.

    Examples
    --------
    Basic usage with tensor:

    >>> from spectre.stats import RunningMean, loader
    >>> import torch
    >>>
    >>> X = torch.randn(10000, 50)
    >>> stat = RunningMean()
    >>>
    >>> for batch in loader(stat, X, batch_size=100):
    ...     stat.add(batch)
    >>>
    >>> mean = stat.mean()  # stat is now on CPU

    Using with DatasetWeighted for weighted statistics:

    >>> from spectre.data import DatasetWeighted
    >>> X = torch.randn(10000, 50)
    >>> weights = torch.rand(10000)
    >>> dataset = DatasetWeighted(X, weights=weights)
    >>> stat = RunningMean()
    >>>
    >>> for batch in loader(stat, dataset, batch_size=100):
    ...     stat.add(batch.data, weights=batch.weights)
    >>>
    >>> mean = stat.mean()

    Subsampling for faster computation:

    >>> # Process only 1000 samples instead of full dataset
    >>> stat = RunningMean()
    >>> for batch in loader(stat, X, batch_size=100, sample_size=1000):
    ...     stat.add(batch)
    >>> mean = stat.mean()

    Computing multiple statistics together:

    >>> from spectre.stats import RunningStat, RunningMean, RunningVariance
    >>>
    >>> stats = RunningStat(mean=RunningMean(), var=RunningVariance())
    >>> for batch in loader(stats, X, batch_size=100):
    ...     stats.add(batch)
    >>>
    >>> mean = stats.mean.mean()
    >>> variance = stats.var.variance()

    Notes
    -----
    The loader automatically moves stat to CPU after processing to avoid
    memory leaks when working with GPU tensors.

    When using DatasetWeighted, the loader yields Batch objects with data,
    weights, and optional target. Use `batch.data` and `batch.weights` when
    calling `stat.add()`.
    """
    if not isinstance(stat, Stats):
        raise TypeError(f"Expected `stat` to be of a `Stats` object, got {type(stat)}.")

    args = {}
    for k in ["sample_size"]:
        if k in kwargs:
            args[k] = kwargs[k]

    stat_loader = make_loader(dataset, **kwargs)

    def wrapped_loader():
        yield from stat_loader
        stat.to(device="cpu")

    return wrapped_loader()

make_loader(dataset: Union[torch.Tensor, DatasetWeighted], sample_size: int | None = None, batch_size: int = 1, sampler: Optional[Sampler] = None, random_seed: int | None = None, **kwargs) -> DataLoader #

Create DataLoader with optional subsampling support.

Helper function for creating DataLoaders with deterministic or random subsampling. Supports deferred dataset loading via callables. Handles both raw tensors and DatasetWeighted instances.

PARAMETER DESCRIPTION
dataset

Data tensor, DatasetWeighted instance, or callable that returns one of these. First dimension is batch/sample dimension.

TYPE: Tensor | DatasetWeighted | Callable

sample_size

Number of samples to process. If None, uses entire dataset.

TYPE: int | None, optional, by default None DEFAULT: None

batch_size

Number of samples per batch.

TYPE: int, optional, by default 1 DEFAULT: 1

sampler

Custom sampler for the DataLoader. Cannot be used with sample_size.

TYPE: Sampler | None, optional, by default None DEFAULT: None

random_seed

Seed for random subsampling. If None and sample_size is set, uses first sample_size samples sequentially.

TYPE: int | None, optional, by default None DEFAULT: None

**kwargs

Additional arguments passed to DataLoader (num_workers, pin_memory, etc.).

DEFAULT: {}

RETURNS DESCRIPTION
DataLoader

Configured DataLoader ready for iteration.
RAISES DESCRIPTION
TypeError

If dataset is not a torch.Tensor, DatasetWeighted, or Callable.
AttributeError

If both sampler and sample_size are specified.

Examples:

Basic DataLoader with tensor:

>>> import torch
>>> from spectre.stats.loader import make_loader
>>>
>>> X = torch.randn(1000, 10)
>>> loader = make_loader(X, batch_size=32)
>>> for batch in loader:
...     print(batch.shape)
torch.Size([32, 10])

DataLoader with DatasetWeighted:

>>> from spectre.data import DatasetWeighted
>>> X = torch.randn(1000, 10)
>>> weights = torch.rand(1000)
>>> dataset = DatasetWeighted(X, weights=weights)
>>> loader = make_loader(dataset, batch_size=32)
>>> for batch in loader:
...     print(batch.data.shape, batch.weights.shape)
torch.Size([32, 10]) torch.Size([32])

Sequential subsampling:

>>> # Process first 200 samples only
>>> loader = make_loader(X, sample_size=200, batch_size=32)
>>> sum(len(batch) for batch in loader)
200

Random subsampling with fixed seed:

>>> # Deterministic random sample of 200 samples
>>> loader = make_loader(X, sample_size=200, batch_size=32, random_seed=42)
Notes

If sample_size exceeds dataset length, a warning is issued and sample_size is clamped to dataset length.

Source code in spectre/stats/loader.py 
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
def make_loader(
    dataset: Union[torch.Tensor, DatasetWeighted],
    sample_size: int | None = None,
    batch_size: int = 1,
    sampler: Optional[Sampler] = None,
    random_seed: int | None = None,
    **kwargs,
) -> DataLoader:
    """
    Create DataLoader with optional subsampling support.

    Helper function for creating DataLoaders with deterministic or random
    subsampling. Supports deferred dataset loading via callables. Handles
    both raw tensors and DatasetWeighted instances.

    Parameters
    ----------
    dataset : torch.Tensor | DatasetWeighted | Callable
        Data tensor, DatasetWeighted instance, or callable that returns one
        of these. First dimension is batch/sample dimension.

    sample_size : int | None, optional, by default None
        Number of samples to process. If None, uses entire dataset.

    batch_size : int, optional, by default 1
        Number of samples per batch.

    sampler : Sampler | None, optional, by default None
        Custom sampler for the DataLoader. Cannot be used with sample_size.

    random_seed : int | None, optional, by default None
        Seed for random subsampling. If None and sample_size is set,
        uses first sample_size samples sequentially.

    **kwargs
        Additional arguments passed to DataLoader (num_workers, pin_memory, etc.).

    Returns
    -------
    DataLoader
        Configured DataLoader ready for iteration.

    Raises
    ------
    TypeError
        If dataset is not a torch.Tensor, DatasetWeighted, or Callable.

    AttributeError
        If both sampler and sample_size are specified.

    Examples
    --------
    Basic DataLoader with tensor:

    >>> import torch
    >>> from spectre.stats.loader import make_loader
    >>>
    >>> X = torch.randn(1000, 10)
    >>> loader = make_loader(X, batch_size=32)
    >>> for batch in loader:
    ...     print(batch.shape)
    torch.Size([32, 10])

    DataLoader with DatasetWeighted:

    >>> from spectre.data import DatasetWeighted
    >>> X = torch.randn(1000, 10)
    >>> weights = torch.rand(1000)
    >>> dataset = DatasetWeighted(X, weights=weights)
    >>> loader = make_loader(dataset, batch_size=32)
    >>> for batch in loader:
    ...     print(batch.data.shape, batch.weights.shape)
    torch.Size([32, 10]) torch.Size([32])

    Sequential subsampling:

    >>> # Process first 200 samples only
    >>> loader = make_loader(X, sample_size=200, batch_size=32)
    >>> sum(len(batch) for batch in loader)
    200

    Random subsampling with fixed seed:

    >>> # Deterministic random sample of 200 samples
    >>> loader = make_loader(X, sample_size=200, batch_size=32, random_seed=42)

    Notes
    -----
    If sample_size exceeds dataset length, a warning is issued and
    sample_size is clamped to dataset length.
    """
    # To support deferred dataset loading, support passing a factory that
    # creates the dataset when called.
    if isinstance(dataset, Callable):
        dataset = dataset()  # ty: ignore [call-non-callable]

    if not isinstance(dataset, (torch.Tensor, DatasetWeighted)):
        raise TypeError(
            f"Expected `dataset` to be `torch.Tensor`, `DatasetWeighted`, "
            f"or `Callable`, got {type(dataset)}."
        )

    if sample_size is not None:
        if sampler:
            raise AttributeError("`Sampler` cannot be specified with `sample_size`.")

        if sample_size > len(dataset):
            warn(f"Sample size {sample_size} > dataset size {len(dataset)}.")
            sample_size = len(dataset)

        if random_seed is None:
            sampler = FixedSubsetSampler(list(range(sample_size)))
        else:
            sampler = FixedRandomSubsetSampler(
                dataset, seed=random_seed, end=sample_size
            )

    dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size, **kwargs)

    return dataloader