Fix/SK-1060 | Move mnist data to scaleout public bucket

examples/mnist-pytorch/client/data.py

@@ -1,23 +1,31 @@
 import os
-from math import floor
 
+import numpy as np
+import requests
 import torch
-import torchvision
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 abs_path = os.path.abspath(dir_path)
 
 
 def get_data(out_dir="data"):
-    # Make dir if necessary
-    if not os.path.exists(out_dir):
-        os.mkdir(out_dir)
-
-    # Only download if not already downloaded
-    if not os.path.exists(f"{out_dir}/train"):
-        torchvision.datasets.MNIST(root=f"{out_dir}/train", transform=torchvision.transforms.ToTensor, train=True, download=True)
-    if not os.path.exists(f"{out_dir}/test"):
-        torchvision.datasets.MNIST(root=f"{out_dir}/test", transform=torchvision.transforms.ToTensor, train=False, download=True)
+    # Generate random int between 1 and 10 for split id, set seed for reproducibility
+    split_id = np.random.randint(1, 11)
+
+    if not os.path.exists(f"{abs_path}/{out_dir}/clients/{split_id}"):
+        # create directory for data
+        os.makedirs(f"{abs_path}/{out_dir}/clients/{split_id}")
+
+    # use requests to download the data from url
+    url = f"https://storage.googleapis.com/public-scaleout/mnist-pytorch/data/clients/{split_id}/mnist.pt"
+    # download into out_dir
+    r = requests.get(url)
+    if r.status_code == 200:
+        with open(f"{abs_path}/{out_dir}/clients/{split_id}/mnist.pt", "wb") as f:
+            f.write(r.content)
+        print(f"Downloaded data from {url}")
+    else:
+        print(f"Failed to download data from {url}")
 
 
 def load_data(data_path, is_train=True):
@@ -31,7 +39,19 @@ def load_data(data_path, is_train=True):
     :rtype: tuple
     """
     if data_path is None:
-        data_path = os.environ.get("FEDN_DATA_PATH", abs_path + "/data/clients/1/mnist.pt")
+        split_id = 0
+        for id in range(1, 11):
+            if os.path.exists(f"{abs_path}/data/clients/{id}/mnist.pt"):
+                split_id = id
+                print(f"Found data at {abs_path}/data/clients/{id}/mnist.pt")
+                break
+        print(f"Using split id {split_id}")
+        data_path = f"{abs_path}/data/clients/{split_id}/mnist.pt"
+        data_path = os.environ.get("FEDN_DATA_PATH", data_path)
+        # check if data_path is a file
+        if not os.path.isfile(data_path):
+            print(f"Data file {data_path} not found.")
+            raise FileNotFoundError
 
     data = torch.load(data_path)
 
@@ -48,50 +68,6 @@ def load_data(data_path, is_train=True):
     return X, y
 
 
-def splitset(dataset, parts):
-    n = dataset.shape[0]
-    local_n = floor(n / parts)
-    result = []
-    for i in range(parts):
-        result.append(dataset[i * local_n : (i + 1) * local_n])
-    return result
-
-
-def split(out_dir="data"):
-    n_splits = int(os.environ.get("FEDN_NUM_DATA_SPLITS", 2))
-
-    # Make dir
-    if not os.path.exists(f"{out_dir}/clients"):
-        os.mkdir(f"{out_dir}/clients")
-
-    # Load and convert to dict
-    train_data = torchvision.datasets.MNIST(root=f"{out_dir}/train", transform=torchvision.transforms.ToTensor, train=True)
-    test_data = torchvision.datasets.MNIST(root=f"{out_dir}/test", transform=torchvision.transforms.ToTensor, train=False)
-    data = {
-        "x_train": splitset(train_data.data, n_splits),
-        "y_train": splitset(train_data.targets, n_splits),
-        "x_test": splitset(test_data.data, n_splits),
-        "y_test": splitset(test_data.targets, n_splits),
-    }
-
-    # Make splits
-    for i in range(n_splits):
-        subdir = f"{out_dir}/clients/{str(i+1)}"
-        if not os.path.exists(subdir):
-            os.mkdir(subdir)
-        torch.save(
-            {
-                "x_train": data["x_train"][i],
-                "y_train": data["y_train"][i],
-                "x_test": data["x_test"][i],
-                "y_test": data["y_test"][i],
-            },
-            f"{subdir}/mnist.pt",
-        )
-
-
 if __name__ == "__main__":
     # Prepare data if not already done
-    if not os.path.exists(abs_path + "/data/clients/1"):
-        get_data()
-        split()
+    get_data()

examples/mnist-pytorch/client/python_env.yaml

@@ -8,8 +8,6 @@ dependencies:
   - torch==2.4.1; (sys_platform == "darwin" and platform_machine == "arm64") or (sys_platform == "win32" or sys_platform == "win64" or sys_platform == "linux")
   # PyTorch macOS x86 builds deprecation
   - torch==2.2.2; sys_platform == "darwin" and platform_machine == "x86_64"
-  - torchvision==0.19.1; (sys_platform == "darwin" and platform_machine == "arm64") or (sys_platform == "win32" or sys_platform == "win64" or sys_platform == "linux")
-  - torchvision==0.17.2; sys_platform == "darwin" and platform_machine == "x86_64"
   - numpy==2.0.2; (sys_platform == "darwin" and platform_machine == "arm64" and python_version >= "3.9") or (sys_platform == "win32" or sys_platform == "win64" or sys_platform == "linux" and python_version >= "3.9")
   - numpy==1.26.4; (sys_platform == "darwin" and platform_machine == "x86_64" and python_version >= "3.9")
   - numpy==1.24.4; python_version == "3.8"