File size: 1,486 Bytes
89434f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | import torch
from torch.utils.data import Dataset, Subset, TensorDataset, random_split
__all__ = [
"make_train_val_indices",
"make_train_val_subsets",
"select_sample_indices",
]
def make_train_val_indices(
n_items: int,
train_fraction: float = 0.8,
seed: int = 42,
) -> tuple[list[int], list[int]]:
if not 0 < train_fraction < 1:
raise ValueError(f"train_fraction must be in (0, 1), got {train_fraction}")
n_train = int(train_fraction * n_items)
generator = torch.Generator().manual_seed(seed)
index_dataset = TensorDataset(torch.arange(n_items))
split_subsets = random_split(
index_dataset,
[n_train, n_items - n_train],
generator=generator,
)
return list(split_subsets[0].indices), list(split_subsets[1].indices)
def make_train_val_subsets(
dataset: Dataset,
train_fraction: float = 0.8,
seed: int = 42,
) -> tuple[Subset, Subset]:
train_indices, val_indices = make_train_val_indices(
len(dataset), # type: ignore[arg-type]
train_fraction=train_fraction,
seed=seed,
)
return Subset(dataset, train_indices), Subset(dataset, val_indices)
def select_sample_indices(n_items: int, n_use: int, seed: int = 42) -> list[int]:
if n_use >= n_items:
return list(range(n_items))
generator = torch.Generator().manual_seed(seed)
chosen = torch.randperm(n_items, generator=generator)[:n_use].tolist()
chosen.sort()
return chosen
|