nemos.batching.DataLoader

class nemos.batching.DataLoader(*args, **kwargs)

Bases: Protocol

Protocol for data loaders that stream batches.

The protocol itself allows batches as tuples of any length, but note that GLM.stochastic_fit expects (X, y) pairs. The variadic batch format is used on the solver-level via AbstractSolver.stochastic_run.

Requirements:

• Must be re-iterable: calling __iter__() must return a fresh iterator each time. This is required for n_passes > 1 and because SVRG’s full gradient computation iterates through the data an additional time per pass.

• sample_batch() should be cheap and deterministic (e.g., the first batch that contains valid data).

• Batches should have consistent, non-zero sizes. Note that the solver’s update method will be recompiled for each unique batch size. This usually means just 2 compilations, as the last batch is almost always of a different size unless the number of samples is divisible by the batch size.

Attributes

n_samples

Total number of samples in the dataset.

__init__(*args, **kwargs)

Methods

__init__(*args, **kwargs)
sample_batch()	Return a single batch for initialization purposes.

__iter__()

Iterate over tuples containing input and output data, e.g. (X_batch, y_batch).

Must return a fresh iterator each call (re-iterable).

Return type:

Iterator[tuple[Any, ...]]

property n_samples: int

Total number of samples in the dataset.

sample_batch()

Return a single batch for initialization purposes.

Should be cheap/cached and deterministic (ignore shuffle setting). Typically returns the first batch that contains valid (non-NaN/Inf) data.

Return type:

tuple[Any, ...]