From d6a00aab979a50394cd4d5423f7304bb7819ce25 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 08:26:40 -0500 Subject: fix: match published baseline protocols --- sdil/data.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) (limited to 'sdil/data.py') 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) -- cgit v1.2.3