Split dataset function for PyTorch

train_test_split.py

from typing import Optional, Tuple

import torch
from torch.utils.data import Dataset, random_split


def train_test_split(
    dataset: Dataset,
    test_ratio: float,
    seed: Optional[int] = None,
) -> Tuple[Dataset, Dataset]:
    """Splits a dataset into random train and test subsets.

    Args:
        dataset (Dataset): dataset.
        test_ratio (float): test proportion (between 0 and 1).
        seed (int, optional): seed. Defaults to None.

    Returns:
        Tuple[Dataset, Dataset]: train and test datasets.
    """
    # Define generator
    generator = torch.Generator()
    if seed is not None:
        generator.manual_seed(seed)

    # Define lengths of subsets
    train_ratio = 1 - test_ratio
    train_size = int(train_ratio * len(dataset))
    test_size = len(dataset) - train_size
    lengths = [train_size, test_size]

    # Split
    train_dataset, test_dataset = random_split(dataset, lengths, generator)

    return train_dataset, test_dataset