kosmos.ml.dataloader

Functions

make_train_test_dataloaders(dataset: kosmos.ml.datasets.dataset.SLDataset, train_split: float, batch_size: int, num_train_subsets: int = 1) tuple[list[torch.utils.data.DataLoader], torch.utils.data.DataLoader]

Split a dataset into training and test subsets and wrap them in DataLoaders.

The training loaders shuffle their subset each epoch; the test loader does not.

Parameters:
  • dataset – The dataset to be split.

  • train_split – Fraction of dataset for training. Test split is 1 - train_split.

  • batch_size – Number of samples per batch in both loaders.

  • num_train_subsets – Number of partitions to split the training subset into. Defaults to 1.

Returns:

A tuple containing:

  • list[DataLoader]: DataLoaders for the partitions of the training subset.

  • DataLoader: DataLoader for the test subset.

Return type:

tuple[list[DataLoader], DataLoader]