diff options
Diffstat (limited to 'sdil/data.py')
| -rw-r--r-- | sdil/data.py | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/sdil/data.py b/sdil/data.py index f5b4928..8253d84 100644 --- a/sdil/data.py +++ b/sdil/data.py @@ -85,7 +85,8 @@ def _load_cifar(data_dir, n_classes=10): torch.from_numpy(norm(xte)), torch.from_numpy(yte)) -def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True): +def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", preload=True, + shuffle_train=True, train_limit=None): if name == "cifar10": xtr, ytr, xte, yte = _load_cifar(data_dir) n_in = 3072 @@ -94,9 +95,13 @@ def get_dataset(name="mnist", batch_size=128, data_dir=DATA_DIR, device="cpu", p xtr, ytr = _load_split(subdir, True, mean, std, data_dir) xte, yte = _load_split(subdir, False, mean, std, data_dir) n_in = 784 + if train_limit is not None: + if train_limit <= 0 or train_limit > xtr.shape[0]: + raise ValueError(f"train_limit must be in [1, {xtr.shape[0]}], got {train_limit}") + xtr, ytr = xtr[:train_limit], ytr[:train_limit] xtr, ytr = xtr.to(device), ytr.to(device) xte, yte = xte.to(device), yte.to(device) - return (_FastLoader(xtr, ytr, batch_size, True), + return (_FastLoader(xtr, ytr, batch_size, shuffle_train), _FastLoader(xte, yte, 1000, False), n_in, 10) |