fix: match published baseline protocols

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)