ty checks passed for examples/callbacks and examples/gradient_leakage_attacks.

Author: baochunli

examples/callbacks/callback_examples.py

@@ -6,6 +6,14 @@
 from plato.callbacks.server import ServerCallback
 from plato.callbacks.trainer import TrainerCallback
 
+__all__ = [
+    "argumentClientCallback",
+    "dynamicClientCallback",
+    "argumentServerCallback",
+    "dynamicServerCallback",
+    "customTrainerCallback",
+]
+
 
 class argumentClientCallback(ClientCallback):
     def on_inbound_received(self, client, inbound_processor):

examples/callbacks/callbacks.py

@@ -2,7 +2,13 @@
 This example shows how to use callbacks to customize server, client, and trainer.
 """
 
-from callback_examples import *
+from callback_examples import (
+    argumentClientCallback,
+    argumentServerCallback,
+    customTrainerCallback,
+    dynamicClientCallback,
+    dynamicServerCallback,
+)
 
 from plato.clients import simple
 from plato.servers import fedavg

examples/client_selection/afl/afl_server.py

@@ -13,11 +13,11 @@
 
 from __future__ import annotations
 
+from afl_selection_strategy import AFLSelectionStrategy
+
 from plato.config import Config
 from plato.servers import fedavg
 
-from afl_selection_strategy import AFLSelectionStrategy
-
 
 class Server(fedavg.Server):
     """An AFL server configured with the strategy-based client selection API."""

examples/client_selection/pisces/pisces_aggregation_strategy.py

@@ -9,6 +9,7 @@
 import asyncio
 import logging
 from types import SimpleNamespace
+
 import numpy as np
 
 from plato.config import Config

examples/composable_trainer/composable_trainer.py

@@ -216,7 +216,9 @@ def setup(self, context: TrainingContext):
         def compute_loss(self, outputs, labels, context):
             """Compute weighted cross-entropy loss + L2 regularization."""
             if self._criterion is None:
-                raise RuntimeError("Loss criterion is not initialized. Call setup first.")
+                raise RuntimeError(
+                    "Loss criterion is not initialized. Call setup first."
+                )
             # Base cross-entropy loss with class weights
             ce_loss = self._criterion(outputs, labels)
 

examples/gradient_leakage_attacks/defense/GradDefense/compensate.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 import torch
+from numpy.typing import ArrayLike, NDArray
 
 
 def get_factor(num):
@@ -45,9 +46,10 @@ def get_matrix_size(total_params_num: int, q: float):
     return gradients_matrix_v, gradients_matrix_w, real_q
 
 
-def get_covariance_matrix(matrix):
+def get_covariance_matrix(matrix: ArrayLike) -> NDArray[np.float64]:
     """Calculate covariance matrix."""
-    return np.cov(matrix, rowvar=0)
+    matrix_array: NDArray[np.float64] = np.asarray(matrix, dtype=np.float64)
+    return np.cov(matrix_array, rowvar=False)
 
 
 def denoise(gradients: list, scale: float, q: float):

examples/gradient_leakage_attacks/defense/GradDefense/dataloader.py

@@ -1,26 +1,41 @@
 """
-Dataloader of GradDefense
+Dataloader of GradDefense.
 
 Reference:
 Wang et al., "Protect Privacy from Gradient Leakage Attack in Federated Learning," INFOCOM 2022.
 https://github.com/wangjunxiao/GradDefense
 """
 
+from __future__ import annotations
+
+from typing import Any, Protocol, Sequence, cast
+
 import numpy as np
 from torch.utils.data import Subset
 from torch.utils.data.dataloader import DataLoader
 from torch.utils.data.dataset import Dataset
 
+
+class LabeledDataset(Protocol):
+    """Minimal protocol for datasets with labeled samples."""
+
+    classes: Sequence[Any]
+
+    def __len__(self) -> int: ...
+
+    def __getitem__(self, index: int) -> tuple[Any, Any]: ...
+
+
 DEFAULT_NUM_WORKERS = 8
 ROOTSET_PER_CLASS = 5
 ROOTSET_SIZE = 50
 
 
 def extract_root_set(
-    dataset: Dataset,
+    dataset: LabeledDataset,
     sample_per_class: int = ROOTSET_PER_CLASS,
-    seed: int = None,
-):
+    seed: int | None = None,
+) -> tuple[list[int], dict[int, list[int]]]:
     """Extract root dataset."""
     num_classes = len(dataset.classes)
     class2sample = {i: [] for i in range(num_classes)}
@@ -39,10 +54,10 @@ def extract_root_set(
     return select_indices, class2sample
 
 
-def get_root_set_loader(trainset):
+def get_root_set_loader(trainset: LabeledDataset) -> DataLoader[Any]:
     """Obtain root dataset loader."""
     rootset_indices, __ = extract_root_set(trainset)
-    root_set = Subset(trainset, rootset_indices)
+    root_set = Subset(cast(Dataset[Any], trainset), rootset_indices)
     root_dataloader = DataLoader(
         root_set, batch_size=len(root_set), num_workers=DEFAULT_NUM_WORKERS
     )

examples/gradient_leakage_attacks/defense/GradDefense/sensitivity.py

@@ -1,34 +1,38 @@
 """
-Sensitivity computation of GradDefense
+Sensitivity computation of GradDefense.
 
 Reference:
 Wang et al., "Protect Privacy from Gradient Leakage Attack in Federated Learning," INFOCOM 2022.
 https://github.com/wangjunxiao/GradDefense
 """
 
+from __future__ import annotations
+
 import torch
 import torch.nn as nn
 from torch.utils.data.dataloader import DataLoader
 
 
 def compute_sens(
     model: nn.Module,
-    rootset_loader: DataLoader,
+    rootset_loader: DataLoader[tuple[torch.Tensor, torch.Tensor]],
     device: torch.device,
-    loss_fn=nn.CrossEntropyLoss(),
-):
+    loss_fn: nn.Module | None = None,
+) -> list[float]:
     """Compute sensitivity."""
     x, y = next(iter(rootset_loader))
 
     x = x.to(device).requires_grad_()
     y = y.to(device)
     model = model.to(device)
 
+    if loss_fn is None:
+        loss_fn = nn.CrossEntropyLoss()
+
     # Compute prediction and loss
-    try:
-        pred, _ = model(x)
-    except:
-        pred = model(x)
+    pred = model(x)
+    if isinstance(pred, tuple):
+        pred = pred[0]
 
     loss = loss_fn(pred, y)
     # Backward propagation
@@ -49,11 +53,11 @@ def compute_sens(
 
     sensitivity = []
     for layer_vjp in vector_jacobian_products:
-        f_norm_sum = 0
+        sum_norms = torch.zeros((), device=layer_vjp.device)
         for sample_vjp in layer_vjp:
             # Sample-wise Frobenius norm
-            f_norm_sum += torch.norm(sample_vjp)
-        f_norm = f_norm_sum / len(layer_vjp)
-        sensitivity.append(f_norm.cpu().numpy())
+            sum_norms = sum_norms + torch.norm(sample_vjp)
+        f_norm = sum_norms / layer_vjp.shape[0]
+        sensitivity.append(float(f_norm.detach().cpu()))
 
     return sensitivity

examples/gradient_leakage_attacks/dlg_model.py

@@ -1,33 +1,32 @@
-"""Obtaining models adapted from existing work's implementations.
+"""
+Obtaining models adapted from existing work's implementations.
 
 An extra return object named `feature` is added in each model's forward function,
 which will be used in the defense Soteria.
 """
 
-from typing import Union
+from __future__ import annotations
+
+from typing import Callable, Optional
 
-from nn import (
-    lenet,
-    resnet,
-)
+import torch.nn as nn
+from nn import lenet, resnet
 
 from plato.config import Config
 
 
-def get(**kwargs: str | dict):
-    """Get the model with the provided name."""
-    model_name = (
-        kwargs["model_name"] if "model_name" in kwargs else Config().trainer.model_name
-    )
+def get(model_name: str | None = None) -> Optional[Callable[[], nn.Module]]:
+    """Get the model constructor with the provided name."""
+    resolved_name = model_name or Config().trainer.model_name
 
-    if model_name == "lenet":
+    if resolved_name == "lenet":
         return lenet.Model
 
-    if model_name.split("_")[0] == "resnet":
-        return resnet.get(model_name=model_name)
+    if resolved_name.startswith("resnet_"):
+        return resnet.get(model_name=resolved_name)
 
     # Set up model through plato's model library
     if Config().trainer.model_type == "vit":
         return None
 
-    raise ValueError(f"No such model: {model_name}")
+    raise ValueError(f"No such model: {resolved_name}")