torch_concepts.data.CUBDataModule¶

class CUBDataModule(root: str | None = None, image_size: int = 224, 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, concept_subset: list | None = None, label_descriptions: dict | None = None, workers: int = 0, **kwargs)[source]¶

DataModule for CUB-200-2011 (Caltech-UCSD Birds).

Handles data loading, splitting, and batching for the CUB-200-2011 dataset with support for concept-based learning. CUB-200-2011 provides official train / val / test splits via the Koh et al. pre-processed pickle files, so NativeSplitter is used by default.

Note

CUB-200-2011 must be manually downloaded before use. See CUBDataset for instructions.

Parameters:

root (str, optional) – Root directory containing class_attr_data_10/ and CUB_200_2011/. Default: None (auto-creates ./data/CUB200).
image_size (int, optional) – Side length (px) to resize images to. Default: 224.
splitter (Splitter, optional) – Splitting strategy. Default: NativeSplitter() (uses the official train / val / test splits from the pickle files).
batch_size (int, optional) – Number of samples per batch. Default: 512.
backbone (BackboneType, optional) – Backbone model for feature extraction (e.g. 'resnet50'). Default: None.
precompute_embs (bool, optional) – Whether to precompute and cache backbone embeddings. Default: True.
force_recompute (bool, optional) – Recompute embeddings even if a cache exists. Default: False.
concept_subset (list of str, optional) – Subset of concept names to retain. Default: None (all 113).
label_descriptions (dict, optional) – Mapping from concept name to human-readable description.
workers (int, optional) – Number of data-loading worker processes. Default: 0.

Examples

>>> from torch_concepts.data import CUBDataModule
>>>
>>> dm = CUBDataModule(
...     root="./data/CUB200",
...     backbone="resnet50",
...     precompute_embs=True,
...     batch_size=64,
... )
>>> dm.setup()
>>> train_loader = dm.train_dataloader()

See also

CUBDataset: The underlying dataset class.
ConceptDataModule: Parent class with common datamodule functionality.

__init__(root: str | None = None, image_size: int = 224, 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, 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, image_size, splitter, ...])
`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.