torch_concepts.data.DSpritesRegressionDataModule

class DSpritesRegressionDataModule(root: str | None = None, formulas: Dict[str, str] | None = None, seed: int = 42, generation_seed: int = 42, splitter: Splitter = RandomSplitter(train_size=None, val_size=None, test_size=None), val_size: int | float = 0.1, test_size: int | float = 0.2, batch_size: int = 512, backbone: str | Callable[[Tensor], Tensor] | None = None, precompute_embs: bool = True, force_recompute: bool = False, concept_subset: list | None = None, label_descriptions: dict | None = None, workers: int = 0, **kwargs)[source]

DataModule for DSprites regression dataset with concept-based learning support.

Handles data loading, splitting, and batching for the DSprites dataset with sympy formula-based regression targets. Supports precomputing backbone embeddings and flexible train/val/test splitting strategies.

Parameters:

  • root (str, optional) – Root directory for caching. Default: None (./data/dsprites_regression).

  • formulas (dict, optional) – Mapping from shape name to sympy formula string. Default: None (uses built-in defaults).

  • seed (int, optional) – Random seed for the train/val/test split. Default: 42

  • generation_seed (int, optional) – Random seed for data generation. Default: 42

  • splitter (Splitter, optional) – Splitting strategy. Default: RandomSplitter()

  • val_size (int or float, optional) – Validation set size. Default: 0.1

  • test_size (int or float, optional) – Test set size. Default: 0.2

  • batch_size (int, optional) – Number of samples per batch. Default: 512

  • backbone (BackboneType, optional) – Backbone model for feature extraction. Default: None

  • precompute_embs (bool, optional) – Whether to precompute and cache backbone embeddings. Default: True

  • force_recompute (bool, optional) – If True, recompute embeddings even if cached. Default: False

  • concept_subset (list of str, optional) – Subset of concept names to retain after loading. Default: None

  • label_descriptions (dict, optional) – Optional dict mapping concept names to descriptions.

  • workers (int, optional) – Number of data loading workers. Default: 0

__init__(root: str | None = None, formulas: Dict[str, str] | None = None, seed: int = 42, generation_seed: int = 42, splitter: Splitter = RandomSplitter(train_size=None, val_size=None, test_size=None), val_size: int | float = 0.1, test_size: int | float = 0.2, batch_size: int = 512, backbone: str | Callable[[Tensor], Tensor] | None = None, precompute_embs: bool = True, force_recompute: bool = False, concept_subset: list | None = None, label_descriptions: dict | None = None, workers: int = 0, **kwargs)[source]
prepare_data_per_node

If True, each LOCAL_RANK=0 will call prepare data. Otherwise only NODE_RANK=0, LOCAL_RANK=0 will prepare data.

allow_zero_length_dataloader_with_multiple_devices

If True, dataloader with zero length within local rank is allowed. Default value is False.

Methods

__init__([root, formulas, seed, ...])

from_datasets([train_dataset, val_dataset, ...])

Create an instance from torch.utils.data.Dataset.

get_dataloader([split, shuffle, batch_size])

Get the DataLoader for a specific split.

load_from_checkpoint(checkpoint_path[, ...])

Primary way of loading a datamodule from a checkpoint.

load_state_dict(state_dict)

Called when loading a checkpoint, implement to reload datamodule state given datamodule state_dict.

on_after_batch_transfer(batch, dataloader_idx)

Override to alter or apply batch augmentations to your batch after it is transferred to the device.

on_before_batch_transfer(batch, dataloader_idx)

Override to alter or apply batch augmentations to your batch before it is transferred to the device.

on_exception(exception)

Called when the trainer execution is interrupted by an exception.

predict_dataloader()

An iterable or collection of iterables specifying prediction samples.

prepare_data()

Use this to download and prepare data.

remove_ignored_hparams(ignore_list)

Remove ignored hyperparameters from the stored state.

save_hyperparameters(*args[, ignore, frame, ...])

Save arguments to hparams attribute.

setup([stage, backbone_device, verbose])

Prepare the data for training, validation, or testing.

state_dict()

Called when saving a checkpoint, implement to generate and save datamodule state.

teardown(stage)

Called at the end of fit (train + validate), validate, test, or predict.

test_dataloader([shuffle, batch_size])

Get the test DataLoader.

train_dataloader([shuffle, batch_size])

Get the training DataLoader.

transfer_batch_to_device(batch, device, ...)

Override this hook if your DataLoader returns tensors wrapped in a custom data structure.

val_dataloader([shuffle, batch_size])

Get the validation DataLoader.

Attributes

CHECKPOINT_HYPER_PARAMS_KEY

CHECKPOINT_HYPER_PARAMS_NAME

CHECKPOINT_HYPER_PARAMS_TYPE

backbone

The backbone model wrapper for feature extraction.

hparams

The collection of hyperparameters saved with save_hyperparameters().

hparams_initial

The collection of hyperparameters saved with save_hyperparameters().

n_samples

Total number of samples in the dataset.

name

test_len

Number of samples in the test set.

testset

The test subset.

train_len

Number of samples in the training set.

trainset

The training subset.

val_len

Number of samples in the validation set.

valset

The validation subset.