Source code for archai.datasets.limit_dataset

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

from typing import List, Tuple, Union, Optional

import torch
from torch.utils.data import \
    SubsetRandomSampler, Sampler, Subset, ConcatDataset, Dataset, random_split

[docs]class LimitDataset(Dataset): def __init__(self, dataset, n): self.dataset = dataset self.n = n if hasattr(dataset, 'targets'): self.targets = dataset.targets[:n] def __len__(self): return self.n def __getitem__(self, i): return self.dataset[i]
DatasetLike = Optional[Union[Dataset, Subset, ConcatDataset, LimitDataset]]