scvi.dataloaders.DataSplitter

class scvi.dataloaders.DataSplitter(adata_manager, train_size=None, validation_size=None, shuffle_set_split=True, load_sparse_tensor=False, pin_memory=False, external_indexing=None, **kwargs)

Bases: LightningDataModule

Creates data loaders train_set, validation_set, test_set.

If train_size + validation_set < 1, then test_set is non-empty.

Parameters:

• adata_manager (AnnDataManager) – AnnDataManager object that has been created via setup_anndata.

• train_size (float | None (default: None)) – float, or None (default is None, which is practically 0.9 and potentially adding a small last batch to validation cells)

• validation_size (float | None (default: None)) – float, or None (default is None)

• shuffle_set_split (bool (default: True)) – Whether to shuffle indices before splitting. If False, the val, train, and test set are split in the sequential order of the data according to validation_size and train_size percentages.

• load_sparse_tensor (bool (default: False)) – EXPERIMENTAL If True, loads sparse CSR or CSC arrays in the input dataset as sparse Tensor with the same layout. Can lead to significant speedups in transferring data to GPUs, depending on the sparsity of the data.

• pin_memory (bool (default: False)) – Whether to copy tensors into device-pinned memory before returning them. Passed into AnnDataLoader.

• external_indexing (list[ndarray] | None (default: None)) – A list of data split indices in the order of training, validation, and test sets. Validation and test set are not required and can be left empty.

• **kwargs – Keyword args for data loader. If adata has labeled data, the data loader class is SemiSupervisedDataLoader, else the data loader class is AnnDataLoader.

Examples

>>> adata = scvi.data.synthetic_iid()
>>> scvi.model.SCVI.setup_anndata(adata)
>>> adata_manager = scvi.model.SCVI(adata).adata_manager
>>> splitter = DataSplitter(adata)
>>> splitter.setup()
>>> train_dl = splitter.train_dataloader()

Attributes table

Methods table

on_after_batch_transfer(batch, dataloader_idx)	Converts sparse tensors to dense if necessary.
setup([stage])	Split indices in train/test/val sets.
test_dataloader()	Create a test data loader.
train_dataloader()	Create a train data loader.
transfer_batch_to_device(batch, device, ...)	Densifies sparse tensors before the transfer if the device can't hold them.
val_dataloader()	Create validation data loader.

Attributes

Methods

DataSplitter.on_after_batch_transfer(batch, dataloader_idx)

Converts sparse tensors to dense if necessary.

DataSplitter.setup(stage=None)

Split indices in train/test/val sets.

DataSplitter.test_dataloader()

Create a test data loader.

DataSplitter.train_dataloader()

Create a train data loader.

DataSplitter.transfer_batch_to_device(batch, device, dataloader_idx)

Densifies sparse tensors before the transfer if the device can’t hold them.

MPS has no sparse CSR/CSC tensor support at all (unlike a missing single-op kernel, the whole layout is unimplemented there), so the default transfer would crash trying to move the still-sparse tensor onto the device; densify first in that case.

DataSplitter.val_dataloader()

Create validation data loader.