enable manuseed from config for mainprocess and dataloaders

config.yaml

@@ -11,6 +11,7 @@ OPTIM:
   LR: 0.1
   MOMENTUM: 0.9
   NAME: SGD
+SEED: 0
 SERVER_EPOCHS: 20
 TEST_SUBPROCESS: false
 WB_ENTITY: example_entity

examples/cuda_demo.py

@@ -17,6 +17,9 @@ def test_fedavg():
     args.NUM_PROCESS = 10
     args.SERVER_EPOCHS = 20
 
+    if args.SEED >= 0:
+        torch.manual_seed(args.SEED)
+
     assert torch.cuda.is_available(), "CUDA is not available"
 
     # 1. Prepare Federated Learning DataSets
@@ -27,7 +30,14 @@ def test_fedavg():
     idx_groups = torch.randperm(effective_size).reshape(args.NUM_CLIENT, -1)
     fed_dss = [ClientDataset(org_ds, idx) for idx in idx_groups.tolist()]
 
-    fed_loader = [DataLoader(ds, args.BATCH_SIZE, shuffle=True) for ds in fed_dss]
+    genetors = [
+        torch.Generator().manual_seed(args.SEED + i) if args.SEED >= 0 else None
+        for i in range(args.NUM_CLIENT)
+    ]
+    fed_loader = [
+        DataLoader(ds, args.BATCH_SIZE, shuffle=True, generator=gtr)
+        for ds, gtr in zip(fed_dss, genetors)
+    ]
     test_loader = DataLoader(test_ds, args.BATCH_SIZE * 4)
 
     # 2. Prepare Model and Criterion

examples/fedavg_demo.py

@@ -14,6 +14,8 @@
 def test_fedavg():
     # 0. Prepare necessary arguments
     args = get_config("config.yaml")
+    if args.SEED >= 0:
+        torch.manual_seed(args.SEED)
 
     # 1. Prepare Federated Learning DataSets
     org_ds = MNIST("dataset", train=True, download=True, transform=ToTensor())
@@ -23,7 +25,14 @@ def test_fedavg():
     idx_groups = torch.randperm(effective_size).reshape(args.NUM_CLIENT, -1)
     fed_dss = [ClientDataset(org_ds, idx) for idx in idx_groups.tolist()]
 
-    fed_loader = [DataLoader(ds, args.BATCH_SIZE, shuffle=True) for ds in fed_dss]
+    genetors = [
+        torch.Generator().manual_seed(args.SEED + i) if args.SEED >= 0 else None
+        for i in range(args.NUM_CLIENT)
+    ]
+    fed_loader = [
+        DataLoader(ds, args.BATCH_SIZE, shuffle=True, generator=gtr)
+        for ds, gtr in zip(fed_dss, genetors)
+    ]
     test_loader = DataLoader(test_ds, args.BATCH_SIZE * 4)
 
     # 2. Prepare Model and Criterion

examples/fedprox_demo.py

@@ -15,6 +15,8 @@ def test_fedavg():
     # 0. Prepare necessary arguments
     args = get_config("config.yaml")
     args.PROX_MU = 0.1
+    if args.SEED >= 0:
+        torch.manual_seed(args.SEED)
 
     # 1. Prepare Federated Learning DataSets
     org_ds = MNIST("dataset", train=True, download=True, transform=ToTensor())
@@ -24,8 +26,15 @@ def test_fedavg():
     idx_groups = torch.randperm(effective_size).reshape(args.NUM_CLIENT, -1)
     fed_dss = [ClientDataset(org_ds, idx) for idx in idx_groups.tolist()]
 
-    fed_loader = [DataLoader(ds, batch_size=32, shuffle=True) for ds in fed_dss]
-    test_loader = DataLoader(test_ds, batch_size=32)
+    genetors = [
+        torch.Generator().manual_seed(args.SEED + i) if args.SEED >= 0 else None
+        for i in range(args.NUM_CLIENT)
+    ]
+    fed_loader = [
+        DataLoader(ds, args.BATCH_SIZE, shuffle=True, generator=gtr)
+        for ds, gtr in zip(fed_dss, genetors)
+    ]
+    test_loader = DataLoader(test_ds, args.BATCH_SIZE * 4)
 
     # 2. Prepare Model and Criterion
     classes = 10