✨ feat(datasets): Separate the data preprocessing

The data preprocessings of trainset and testset are defined separately with no overlap.

Add `train()` and `eval()` for `BaseDataset` that for informing FL-bench that which data preprocessing should be adopted.

🐞fix(Elastic): Remove deprecated arg in torch's lr_scheduler

🐞fix(ADCOL): Add definitions of some variables in `__init__()`

data/utils/datasets.py

@@ -22,22 +22,28 @@ def __init__(self) -> None:
         self.targets: torch.Tensor = None
         self.train_data_transform = None
         self.train_target_transform = None
-        self.general_data_transform = None
-        self.general_target_transform = None
-        self.enable_train_transform = True
+        self.test_data_transform = None
+        self.test_target_transform = None
+        self.data_transform = None
+        self.target_transform = None
 
     def __getitem__(self, index):
         data, targets = self.data[index], self.targets[index]
-        if self.enable_train_transform and self.train_data_transform is not None:
-            data = self.train_data_transform(data)
-        if self.enable_train_transform and self.train_target_transform is not None:
-            targets = self.train_target_transform(targets)
-        if self.general_data_transform is not None:
-            data = self.general_data_transform(data)
-        if self.general_target_transform is not None:
-            targets = self.general_target_transform(targets)
+        if self.data_transform is not None:
+            data = self.data_transform(data)
+        if self.target_transform is not None:
+            targets = self.target_transform(targets)
+
         return data, targets
 
+    def train(self):
+        self.data_transform = self.train_data_transform
+        self.target_transform = self.train_target_transform
+
+    def eval(self):
+        self.data_transform = self.test_data_transform
+        self.target_transform = self.test_target_transform
+
     def __len__(self):
         return len(self.targets)
 
@@ -47,8 +53,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ) -> None:
@@ -68,8 +74,8 @@ def __init__(
         self.data = torch.from_numpy(data).float().reshape(-1, 1, 28, 28)
         self.targets = torch.from_numpy(targets).long()
         self.classes = list(range(62))
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -99,8 +105,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ) -> None:
@@ -119,8 +125,8 @@ def __init__(
 
         self.data = torch.from_numpy(data).permute([0, -1, 1, 2]).float()
         self.targets = torch.from_numpy(targets).long()
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
         self.classes = [0, 1]
@@ -131,8 +137,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -146,8 +152,8 @@ def __init__(
         self.targets = (
             torch.Tensor(np.load(root / "raw" / "ydata.npy")).long().squeeze()
         )
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -157,8 +163,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -174,8 +180,8 @@ def __init__(
             torch.Tensor(np.load(root / "raw" / "ydata.npy")).long().squeeze()
         )
         self.classes = [0, 1, 2, 3]
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -185,8 +191,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -203,8 +209,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = list(range(10))
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -214,8 +220,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -232,8 +238,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = list(range(10))
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -243,8 +249,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -258,8 +264,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = train_part.classes
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -269,8 +275,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -284,8 +290,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = train_part.classes
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -295,8 +301,8 @@ def __init__(
         self,
         root,
         args,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -319,8 +325,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = train_part.classes
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -330,8 +336,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -345,8 +351,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = train_part.classes
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -356,8 +362,8 @@ def __init__(
         self,
         root,
         args,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -371,8 +377,8 @@ def __init__(
         self.data = torch.cat([train_data, test_data])
         self.targets = torch.cat([train_targets, test_targets])
         self.classes = train_part.classes
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
         super_class = None
@@ -420,8 +426,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -472,8 +478,8 @@ def __init__(
 
         self.data = torch.load(root / "data.pt")
         self.targets = torch.load(root / "targets.pt")
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -483,8 +489,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ):
@@ -528,8 +534,8 @@ def __init__(
 
         self.data = torch.load(root / "data.pt")
         self.targets = torch.load(root / "targets.pt")
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
@@ -539,8 +545,8 @@ def __init__(
         self,
         root,
         args=None,
-        general_data_transform=None,
-        general_target_transform=None,
+        test_data_transform=None,
+        test_target_transform=None,
         train_data_transform=None,
         train_target_transform=None,
     ) -> None:
@@ -576,22 +582,18 @@ def __init__(
                 transforms.ToTensor(),
             ]
         )
-        self.general_data_transform = general_data_transform
-        self.general_target_transform = general_target_transform
+        self.test_data_transform = test_data_transform
+        self.test_target_transform = test_target_transform
         self.train_data_transform = train_data_transform
         self.train_target_transform = train_target_transform
 
     def __getitem__(self, index):
         data = self.pre_transform(Image.open(self.filename_list[index]).convert("RGB"))
         targets = self.targets[index]
-        if self.enable_train_transform and self.train_data_transform is not None:
-            data = self.train_data_transform(data)
-        if self.enable_train_transform and self.train_target_transform is not None:
-            targets = self.train_target_transform(targets)
-        if self.general_data_transform is not None:
-            data = self.general_data_transform(data)
-        if self.general_target_transform is not None:
-            targets = self.general_target_transform(targets)
+        if self.data_transform is not None:
+            data = self.data_transform(data)
+        if self.target_transform is not None:
+            targets = self.target_transform(targets)
         return data, targets
 
 

src/client/adcol.py

@@ -15,10 +15,12 @@ def __init__(self, model, discriminator, args, logger, device, client_num):
         self.discriminator = discriminator
         self.discriminator.to(self.device)
         self.client_num = client_num
+        self.featrure_list = []
 
     def fit(self):
         self.model.train()
         self.discriminator.eval()
+        self.dataset.train()
         self.featrure_list = []
         for i in range(self.local_epoch):
             for x, y in self.trainloader:

src/client/apfl.py

@@ -53,6 +53,7 @@ def save_state(self):
     def fit(self):
         self.model.train()
         self.local_model.train()
+        self.dataset.train()
         for i in range(self.local_epoch):
             for x, y in self.trainloader:
                 if len(x) <= 1:

src/client/ditto.py

@@ -32,6 +32,7 @@ def save_state(self):
 
     def fit(self):
         self.model.train()
+        self.dataset.train()
         for _ in range(self.local_epoch):
             for x, y in self.trainloader:
                 if len(x) <= 1: