MNISTDataModule#

class torch_uncertainty.datamodules.MNISTDataModule(root, batch_size, eval_batch_size=None, eval_ood=False, eval_shift=False, ood_ds='fashion', num_tta=1, val_split=None, postprocess_set='val', num_workers=1, train_transform=None, test_transform=None, ood_transform=None, basic_augment=True, cutout=None, pin_memory=True, persistent_workers=True)[source]#

DataModule for MNIST.

Parameters:

  • root (str | Path) – Root directory of the datasets.

  • eval_ood (bool) – Whether to evaluate on out-of-distribution data. Defaults to False.

  • eval_shift (bool) – Whether to evaluate on shifted data. Defaults to False.

  • batch_size (int) – Number of samples per batch during training.

  • eval_batch_size (int | None) – Number of samples per batch during evaluation (val and test). Set to batch_size if None. Defaults to None.

  • ood_ds (Literal['fashion', 'notMNIST']) – Which out-of-distribution dataset to use. Defaults to "fashion"; fashion stands for FashionMNIST and notMNIST for notMNIST.

  • val_split (float | None) – Share of samples to use for validation. Defaults to 0.0.

  • num_tta (int) – Number of test-time augmentations (TTA). Defaults to 1 (no TTA).

  • postprocess_set (Literal['val', 'test']) – The post-hoc calibration dataset to use for the post-processing method. Defaults to val.

  • num_workers (int) – Number of workers to use for data loading. Defaults to 1.

  • train_transform (Module | None) – Custom training transform. Defaults to None. If not provided, a default transform is used.

  • test_transform (Module | None) – Custom test transform. Defaults to None. If not provided, a default transform is used.

  • ood_transform (Module | None) – Custom transform for out-of-distribution datasets. Defaults to None. If not provided, a default transform is used.

  • basic_augment (bool) – Whether to apply base augmentations. Defaults to True.

  • cutout (int | None) – Size of cutout to apply to images. Defaults to None.

  • pin_memory (bool) – Whether to pin memory. Defaults to True.

  • persistent_workers (bool) – Whether to use persistent workers. Defaults to True.

prepare_data()[source]#

Download the datasets.

Return type:

None

setup(stage=None)[source]#

Set up the datasets for training, validation, and testing.

Parameters:

stage (Literal["fit", "test"] | None) – Stage of the setup process. Defaults to None.

Return type:

None

test_dataloader()[source]#

Get the test dataloaders for MNIST.

Returns:

Dataloaders of the MNIST test set (in

distribution data), FashionMNIST or NotMNIST test split (out-of-distribution data), and/or MNISTC (shifted data).

Return type:

list[DataLoader]