Update the codebase for gradient leakage attacks

examples/gradient_leakage_attacks/README.md

@@ -23,7 +23,7 @@ Research on gradient leakage attack and defense using the Plato framework.
 - {GC} gradient compression (referred to Soteria's code) 
 ---
 
-# Running Plato with DLG attack
+## Run FL with DLG attack
 Run the program using the following cmd under the directory `examples/gradient_leakage_attacks` for example, which launches the csDLG attack using delta matching against an untrained ResNet18 model in `.eval()` mode on CIFAR100 data:
 ```
 python dlg.py -c untrained_eval_delta.yml
@@ -87,6 +87,9 @@ Try tuning the following hyperparameters in `.yml` configuration files.
 - (for Outpost) `prune_base: 80` — pruning percentage
 - (for Outpost) `noise_base: 0.8` — scale for gaussian noise
 
+### under `trainer`
+
+- `model_type: vit` — needed when using ViTs as the model; specify the ViT model name such as `vit-base-patch16-224-in21k` to `model_name` at the same time
 
 ### under `results`
 
@@ -95,9 +98,5 @@ Try tuning the following hyperparameters in `.yml` configuration files.
 - `trial: [int]` — the trial number to be plotted. Plots the best reconstruction based off MSE if not specified
 
 - `cols/rows: [int]` — specify the number of columns or rows in the plotted figure
-
-## Plot Instructions
-
-Run ```python plot.py -c config_file``` where ```config_file``` is the same one used to run the DLG attack
 
 

examples/gradient_leakage_attacks/defense/GradDefense/compensate.py

@@ -1,10 +1,19 @@
+"""
+Gradient compensating of GradDefense
+
+Reference:
+Wang et al., "Protect Privacy from Gradient Leakage Attack in Federated Learning," INFOCOM 2022.
+https://github.com/wangjunxiao/GradDefense
+"""
+
 import math
 
 import numpy as np
 import torch
 
 
-def getFactor(num):
+def get_factor(num):
+    """Calculate factors."""
     factors = []
     for_times = int(math.sqrt(num))
     for i in range(for_times + 1)[1:]:
@@ -16,10 +25,11 @@ def getFactor(num):
     return factors
 
 
-def getMatrixSize(total_params_num: int, Q: float):
-    gradients_matrix_v = math.sqrt(Q * total_params_num)
+def get_matrix_size(total_params_num: int, q: float):
+    """Calculate matrix size."""
+    gradients_matrix_v = math.sqrt(q * total_params_num)
 
-    for factor in sorted(getFactor(total_params_num)):
+    for factor in sorted(get_factor(total_params_num)):
         if factor >= gradients_matrix_v:
             gradients_matrix_v = factor
             break
@@ -30,84 +40,71 @@ def getMatrixSize(total_params_num: int, Q: float):
     assert isinstance(gradients_matrix_w, int)
     assert gradients_matrix_v * gradients_matrix_w == total_params_num
 
-    real_Q = gradients_matrix_v / gradients_matrix_w
+    real_q = gradients_matrix_v / gradients_matrix_w
 
-    return gradients_matrix_v, gradients_matrix_w, real_Q
+    return gradients_matrix_v, gradients_matrix_w, real_q
 
 
-def getCovarianceMatrix(matrix):
+def get_covariance_matrix(matrix):
+    """Calculate covariance matrix."""
     return np.cov(matrix, rowvar=0)
 
 
-def denoise(gradients: list, scale: float, Q: float):
+def denoise(gradients: list, scale: float, q: float):
+    """Denoise gradients."""
     layer_dims_pool = []
     for layer in gradients:
         layer_dims = list((_ for _ in layer.shape))
         layer_dims_pool.append(layer_dims)
 
-    # print(layer_dims_pool)
-
     layer_params_num_pool = []
     for layer_dims in layer_dims_pool:
         layer_params_num = 1
         for dim in layer_dims:
             layer_params_num *= dim
         layer_params_num_pool.append(layer_params_num)
 
-    # print(layer_params_num_pool)
-
     total_params_num = 0
     for layer_params_num in layer_params_num_pool:
         total_params_num += layer_params_num
 
-    # print(total_params_num)
-
-    gradients_matrix_v, gradients_matrix_w, real_Q = getMatrixSize(
-        total_params_num=total_params_num, Q=Q
+    gradients_matrix_v, gradients_matrix_w, real_q = get_matrix_size(
+        total_params_num=total_params_num, q=q
     )
 
-    # print (gradients_matrix_v, gradients_matrix_w, real_Q)
-
     # Flatten gradients
     gradients_flatten = []
     for layer in gradients:
         layer_flatten = (torch.flatten(layer)).cpu().numpy()
         gradients_flatten.extend(layer_flatten)
 
-    # print(len(gradients_flatten))
-
-    matrix_C = np.array(gradients_flatten).reshape(
+    matrix_c = np.array(gradients_flatten).reshape(
         gradients_matrix_v, gradients_matrix_w
     )
-    covmatrix_CTC = getCovarianceMatrix(matrix_C)
-
-    # print(covmatrix_CTC.shape)
+    covmatrix_ctc = get_covariance_matrix(matrix_c)
 
-    lamda_min = ((1 - 1 / math.sqrt(real_Q)) ** 2) * (scale**2)
-    lamda_max = ((1 + 1 / math.sqrt(real_Q)) ** 2) * (scale**2)
+    lamda_min = ((1 - 1 / math.sqrt(real_q)) ** 2) * (scale**2)
+    lamda_max = ((1 + 1 / math.sqrt(real_q)) ** 2) * (scale**2)
 
-    eigen_vals, eigen_vecs = np.linalg.eig(covmatrix_CTC)
+    eigen_vals, eigen_vecs = np.linalg.eig(covmatrix_ctc)
 
     n_index = []
-    for index in range(len(eigen_vals)):
-        if eigen_vals[index] <= lamda_min or eigen_vals[index] >= lamda_max:
+    for index, eigen_val in enumerate(eigen_vals):
+        if eigen_val <= lamda_min or eigen_val >= lamda_max:
             n_index.append(index)
     n_eigen_vecs = eigen_vecs[:, n_index]
 
-    lowData = np.dot(matrix_C, n_eigen_vecs)
-    highData = np.dot(lowData, n_eigen_vecs.T)
-
-    # print(n_eigen_vecs.shape)
-    # print(lowData.shape)
-    # print(highData.shape)
+    low_data = np.dot(matrix_c, n_eigen_vecs)
+    high_data = np.dot(low_data, n_eigen_vecs.T)
 
-    compensated_gradients_flatten = highData.flatten()
+    compensated_gradients_flatten = high_data.flatten()
 
     gradients_compensated = []
     params_start_indice = 0
     params_end_indice = 0
+
     # Recover gradients
-    for layer_index in range(len(layer_dims_pool)):
+    for layer_index, _ in enumerate(layer_dims_pool):
         if layer_index == 0:
             params_start_indice = 0
             params_end_indice = layer_params_num_pool[layer_index] - 1
@@ -117,9 +114,6 @@ def denoise(gradients: list, scale: float, Q: float):
                 params_start_indice + layer_params_num_pool[layer_index] - 1
             )
 
-        # print(params_start_indice, params_end_indice)
-        # print(layer_dims_pool[layer_index])
-
         layer_gradient = (
             torch.from_numpy(
                 np.array(

examples/gradient_leakage_attacks/defense/GradDefense/dataloader.py

@@ -1,26 +1,31 @@
+"""
+Dataloader of GradDefense
+
+Reference:
+Wang et al., "Protect Privacy from Gradient Leakage Attack in Federated Learning," INFOCOM 2022.
+https://github.com/wangjunxiao/GradDefense
+"""
+
 import numpy as np
-import torch
 from torch.utils.data import Subset
 from torch.utils.data.dataloader import DataLoader
 from torch.utils.data.dataset import Dataset
 
 DEFAULT_NUM_WORKERS = 8
-rootset_per_class = 5
-rootset_size = 50
-
-# TODO: total_num_samples is not used
+ROOTSET_PER_CLASS = 5
+ROOTSET_SIZE = 50
 
 
 def extract_root_set(
     dataset: Dataset,
-    sample_per_class: int = rootset_per_class,
-    total_num_samples: int = rootset_size,
+    sample_per_class: int = ROOTSET_PER_CLASS,
     seed: int = None,
 ):
+    """Extract root dataset."""
     num_classes = len(dataset.classes)
     class2sample = {i: [] for i in range(num_classes)}
     select_indices = []
-    if seed == None:
+    if seed is None:
         index_pool = range(len(dataset))
     else:
         index_pool = np.random.RandomState(seed=seed).permutation(len(dataset))
@@ -35,6 +40,7 @@ def extract_root_set(
 
 
 def get_root_set_loader(trainset):
+    """Obtain root dataset loader."""
     rootset_indices, __ = extract_root_set(trainset)
     root_set = Subset(trainset, rootset_indices)
     root_dataloader = DataLoader(