torch_concepts.data.MNISTArithmeticDataModule¶

class MNISTArithmeticDataModule(root: str | None = None, num_train_samples: int = 10000, num_test_samples: int = 2000, val_size: float = 0.1, img_size: int = 224, seed: int = 42, generation_seed: int = 42, splitter: Splitter = NativeSplitter(train_size=None, val_size=None, test_size=None), batch_size: int = 512, backbone: str | Callable[[Tensor], Tensor] | None = None, precompute_embs: bool = True, force_recompute: bool = False, label_descriptions: dict | None = None, workers: int = 0, **kwargs)[source]¶

DataModule for MNIST Arithmetic dataset with concept-based learning support.

Handles data loading, splitting, and batching for the MNIST arithmetic composite image dataset. Training/validation composites use MNIST train digits while test composites use MNIST test digits, preventing digit-level leakage. Supports precomputing backbone embeddings.

Parameters:

root (str, optional) – Root directory where the dataset is stored or will be generated. Default: None (auto-creates ./data/mnist_arithmetic).
num_train_samples (int, optional) – Number of composite samples from MNIST train. Default: 10000
num_test_samples (int, optional) – Number of composite samples from MNIST test. Default: 2000
val_size (float, optional) – Fraction of MNIST-train composites for validation. Default: 0.1
img_size (int, optional) – Output image size (square). Default: 224
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: NativeSplitter() (uses the native train/val/test mapping built from MNIST splits).
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
workers (int, optional) – Number of data loading workers. Default: 0

Examples

>>> from torch_concepts.data import MNISTArithmeticDataModule
>>>
>>> dm = MNISTArithmeticDataModule(
...     num_train_samples=1000,
...     num_test_samples=200,
...     img_size=224,
...     batch_size=32, seed=42,
... )
>>> dm.setup()
>>> train_loader = dm.train_dataloader()

__init__(root: str | None = None, num_train_samples: int = 10000, num_test_samples: int = 2000, val_size: float = 0.1, img_size: int = 224, seed: int = 42, generation_seed: int = 42, splitter: Splitter = NativeSplitter(train_size=None, val_size=None, test_size=None), batch_size: int = 512, backbone: str | Callable[[Tensor], Tensor] | None = None, precompute_embs: bool = True, force_recompute: bool = False, 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, num_train_samples, ...])
`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.