peal.data.dataloaders.DataloaderMixer

class peal.data.dataloaders.DataloaderMixer

Bases: DataLoader

Dataloader that samples batches from several dataloaders.

Starts with one dataloader; append adds more (e.g. counterfactual datasets during CFKD) with priorities proportional to dataset size. In the default mode every sample() picks one dataloader by multinomial draw over priorities; with train_config.concatenate_batches one batch from each dataloader is concatenated instead (used for class balancing). Exhausted iterators are restarted transparently, and one epoch is train_config.steps_per_epoch batches. Only the attribute interface of DataLoader is reused; DataLoader.__init__ is not called.

Parameters:

• train_config (TrainingConfig) – Provides steps_per_epoch and optionally concatenate_batches.

• initial_dataloader (DataLoader) – First member; its batch_size and dataset are adopted.

• return_src (bool, optional) – If True (and batches are not concatenated) sample returns (batch, dataloader_index).

Variables:

• dataloaders (list of DataLoader)

• iterators (list) – One live iterator per dataloader.

• priorities (numpy.ndarray or None) – Sampling probabilities, None while only one dataloader exists.

• class_balancing_enabled (hints_enabled,)

Methods

append	Add another dataloader and recompute the sampling priorities.
disable_class_balancing	Undo enable_class_balancing.
disable_hints	Disable hint outputs on all member datasets and reset the iterators.
disable_idx	Stop member datasets returning sample indices and reset the iterators.
enable_class_balancing	Replace every plain member dataloader by a per-class concatenating mixer.
enable_hints	Enable hint outputs on all member datasets and reset the iterators.
enable_idx	Make member datasets return sample indices and reset the iterators.
remove_empty_dataloaders	Recursively remove dataloaders with zero-length datasets.
reset	Rebuild every member dataloader and its iterator.
sample	Draw one batch.
update_dataset	Re-point self.dataset at the first member's dataset (arg ignored).

Attributes

return_src

Whether sample also returns the source index (off when concatenating).

__init__(train_config, initial_dataloader, return_src=False)

Register the first dataloader (DataLoader.__init__ is skipped).

update_dataset(dataset)

Re-point self.dataset at the first member’s dataset (arg ignored).

property return_src

Whether sample also returns the source index (off when concatenating).

append(dataloader, priority=1, weight_added_dataloader=None)

Add another dataloader and recompute the sampling priorities.

Parameters:

• dataloader (DataLoader) – Dataloader to add.

• priority (float, optional) – Multiplier on the new dataloader’s dataset size before the priorities are normalized. Defaults to 1.

• weight_added_dataloader (float, optional) – If given, priorities become [1 - w, w] regardless of sizes (only meaningful with exactly two dataloaders).

Notes

With train_config.concatenate_batches all members are reset to half the batch size so concatenated batches keep the original size.

sample()

Draw one batch.

In priority mode one dataloader is chosen by multinomial draw and its next batch returned; in concatenation mode the next batch of every dataloader is fetched and the tensors (also inside nested lists/tuples) are concatenated along dim 0. Exhausted iterators are restarted once; a dataloader that is still empty afterwards raises StopIteration.

Returns:

The batch as produced by the member dataloader(s), or (batch, source_index) when return_src is set.

Return type:

object

reset(batch_size=None)

Rebuild every member dataloader and its iterator.

Parameters:

batch_size (int, optional) – New batch size for the rebuilt DataLoader objects; nested mixers are reset recursively with the same value.

remove_empty_dataloaders()

Recursively remove dataloaders with zero-length datasets.

This method iterates through all nested dataloaders and removes any that have no data (len(dataset) == 0), including nested DataloaderMixers that become empty after recursive cleaning.

enable_hints()

Enable hint outputs on all member datasets and reset the iterators.

disable_hints()

Disable hint outputs on all member datasets and reset the iterators.

enable_idx()

Make member datasets return sample indices and reset the iterators.

disable_idx()

Stop member datasets returning sample indices and reset the iterators.

enable_class_balancing()

Replace every plain member dataloader by a per-class concatenating mixer.

For each class i in dataset.output_size a deep copy of the dataloader restricted to that class is created; the copies are wrapped in a nested DataloaderMixer with concatenate_batches and steps_per_epoch = 200 so each batch holds an equal share of every class. No-op if already enabled.

disable_class_balancing()

Undo enable_class_balancing.

Keeps only the first plain member dataloader found (with its class restriction removed) and drops the priorities. No-op if not enabled.