refactor: use alexnet to replace efficientnet

examples/advanced-pytorch/README.md

@@ -69,3 +69,11 @@ but this can be changed by removing the `--toy` argument in the script. You can
 The `run.sh` script starts processes in the background so that you don't have to open eleven terminal windows. If you experiment with the code example and something goes wrong, simply using `CTRL + C` on Linux (or `CMD + C` on macOS) wouldn't normally kill all these processes, which is why the script ends with `trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT` and `wait`. This simply allows you to stop the experiment using `CTRL + C` (or `CMD + C`). If you change the script and anything goes wrong you can still use `killall python` (or `killall python3`) to kill all background processes (or a more specific command if you have other Python processes running that you don't want to kill).
 
 You can also manually run `python3 server.py` and `python3 client.py --client-id <ID>` for as many clients as you want but you have to make sure that each command is run in a different terminal window (or a different computer on the network).
+
+
+## About Differential Privacy
+If you want to achieve differential privacy, please use `alexnet` model, because efficientnet is particularly affected by noise. 
+
+If you add a little noise when using `efficientnet`, the loss will be Nan. This is issue:https://github.com/adap/flower/issues/2342
+
+Therefore, if you want to achieve differential privacy, please use `alexnet`

examples/advanced-pytorch/alexnet.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class AlexNet(nn.Module):  # lr = 0.01
+    def __init__(self, class_num=10):
+        super(AlexNet, self).__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(64, 192, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(192, 384, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+
+        self.classifier = nn.Sequential(
+            nn.Dropout(0.75),
+            nn.Linear(256 * 4 * 4, 4096),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.75),
+            nn.Linear(4096, 4096),
+            nn.ReLU(inplace=True),
+            nn.Linear(4096, class_num),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = x.view(x.size(0), 256 * 4 * 4)
+        x = self.classifier(x)
+        return x

examples/advanced-pytorch/client.py

@@ -17,26 +17,30 @@ def __init__(
         trainset: datasets.Dataset,
         testset: datasets.Dataset,
         device: torch.device,
+        model_str: str,
         validation_split: int = 0.1,
     ):
         self.device = device
         self.trainset = trainset
         self.testset = testset
         self.validation_split = validation_split
+        if model_str == "alexnet":
+            self.model = utils.load_alexnet(classes=10)
+        else:
+            self.model = utils.load_efficientnet(classes=10)
 
     def set_parameters(self, parameters):
-        """Loads a efficientnet model and replaces it parameters with the ones given."""
-        model = utils.load_efficientnet(classes=10)
-        params_dict = zip(model.state_dict().keys(), parameters)
+        """Loads a alexnet or efficientnet model and replaces it parameters with the ones given."""
+
+        params_dict = zip(self.model.state_dict().keys(), parameters)
         state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
-        model.load_state_dict(state_dict, strict=True)
-        return model
+        self.model.load_state_dict(state_dict, strict=True)
 
     def fit(self, parameters, config):
         """Train parameters on the locally held training set."""
 
         # Update local model parameters
-        model = self.set_parameters(parameters)
+        self.set_parameters(parameters)
 
         # Get hyperparameters for this round
         batch_size: int = config["batch_size"]
@@ -49,25 +53,25 @@ def fit(self, parameters, config):
         train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
         val_loader = DataLoader(valset, batch_size=batch_size)
 
-        results = utils.train(model, train_loader, val_loader, epochs, self.device)
+        results = utils.train(self.model, train_loader, val_loader, epochs, self.device)
 
-        parameters_prime = utils.get_model_params(model)
+        parameters_prime = utils.get_model_params(self.model)
         num_examples_train = len(trainset)
 
         return parameters_prime, num_examples_train, results
 
     def evaluate(self, parameters, config):
         """Evaluate parameters on the locally held test set."""
         # Update local model parameters
-        model = self.set_parameters(parameters)
+        self.set_parameters(parameters)
 
         # Get config values
         steps: int = config["val_steps"]
 
         # Evaluate global model parameters on the local test data and return results
         testloader = DataLoader(self.testset, batch_size=16)
 
-        loss, accuracy = utils.test(model, testloader, steps, self.device)
+        loss, accuracy = utils.test(self.model, testloader, steps, self.device)
         return float(loss), len(self.testset), {"accuracy": float(accuracy)}
 
 
@@ -121,6 +125,14 @@ def main() -> None:
         required=False,
         help="Set to true to use GPU. Default: False",
     )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="efficientnet",
+        choices=['efficientnet', 'alexnet'],
+        help="Use either Efficientnet or Alexnet models. \
+             If you want to achieve differential privacy, please use the Alexnet model",
+    )
 
     args = parser.parse_args()
 
@@ -138,7 +150,7 @@ def main() -> None:
             trainset = trainset.select(range(10))
             testset = testset.select(range(10))
         # Start Flower client
-        client = CifarClient(trainset, testset, device)
+        client = CifarClient(trainset, testset, device, args.model)
 
         fl.client.start_numpy_client(server_address="127.0.0.1:8080", client=client)
 

examples/advanced-pytorch/server.py

@@ -51,9 +51,9 @@ def get_evaluate_fn(model: torch.nn.Module, toy: bool):
 
     # The `evaluate` function will be called after every round
     def evaluate(
-        server_round: int,
-        parameters: fl.common.NDArrays,
-        config: Dict[str, fl.common.Scalar],
+            server_round: int,
+            parameters: fl.common.NDArrays,
+            config: Dict[str, fl.common.Scalar],
     ) -> Optional[Tuple[float, Dict[str, fl.common.Scalar]]]:
         # Update model with the latest parameters
         params_dict = zip(model.state_dict().keys(), parameters)
@@ -80,17 +80,28 @@ def main():
         help="Set to true to use only 10 datasamples for validation. \
             Useful for testing purposes. Default: False",
     )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="efficientnet",
+        choices=['efficientnet', 'alexnet'],
+        help="Use either Efficientnet or Alexnet models. \
+             If you want to achieve differential privacy, please use the Alexnet model",
+    )
 
     args = parser.parse_args()
 
-    model = utils.load_efficientnet(classes=10)
+    if args.model == "alexnet":
+        model = utils.load_alexnet(classes=10)
+    else:
+        model = utils.load_efficientnet(classes=10)
 
     model_parameters = [val.cpu().numpy() for _, val in model.state_dict().items()]
 
     # Create strategy
     strategy = fl.server.strategy.FedAvg(
-        fraction_fit=0.2,
-        fraction_evaluate=0.2,
+        fraction_fit=1.0,
+        fraction_evaluate=1.0,
         min_fit_clients=2,
         min_evaluate_clients=2,
         min_available_clients=10,

examples/advanced-pytorch/utils.py

@@ -1,11 +1,12 @@
 import torch
 from torchvision.transforms import Compose, ToTensor, Normalize, Resize, CenterCrop
 from torch.utils.data import DataLoader
-
+from alexnet import AlexNet
 import warnings
 
 from flwr_datasets import FederatedDataset
 
+
 warnings.filterwarnings("ignore")
 
 
@@ -129,3 +130,7 @@ def load_efficientnet(entrypoint: str = "nvidia_efficientnet_b0", classes: int =
 def get_model_params(model):
     """Returns a model's parameters."""
     return [val.cpu().numpy() for _, val in model.state_dict().items()]
+
+
+def load_alexnet(classes):
+    return AlexNet(classes)