Formatted all modules according to pep8

Author: LongPham7

adversarial_attack.py

@@ -12,6 +12,7 @@
 This module contains classes for adversarial attacks.
 """
 
+
 class FGSM:
     """
     Class for the fast gradient sign method (FGSM).
@@ -23,10 +24,12 @@ def __init__(self, model, loss_criterion, norm, batch_size=128):
         self.pytorch_model = wrapModel(model, loss_criterion)
         self.norm = norm
         self.batch_size = batch_size
-        self.attack = FastGradientMethod(self.pytorch_model, batch_size=batch_size)
-        
+        self.attack = FastGradientMethod(
+            self.pytorch_model, batch_size=batch_size)
+
         # Use GPU for computation if it is available
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu")
 
     def generatePerturbation(self, data, budget, minimal=False):
         """
@@ -46,12 +49,14 @@ def generatePerturbation(self, data, budget, minimal=False):
         """
 
         images, _ = data
-        images_adv = self.attack.generate(x=images.cpu().numpy(), norm=self.norm, eps=budget, minimal=minimal, eps_step=budget / 50, eps_max=budget, batch_size=self.batch_size)
+        images_adv = self.attack.generate(x=images.cpu().numpy(
+        ), norm=self.norm, eps=budget, minimal=minimal, eps_step=budget / 50, eps_max=budget, batch_size=self.batch_size)
         images_adv = torch.from_numpy(images_adv)
 
-        # The output to be returned should be loaded on an appropriate device. 
+        # The output to be returned should be loaded on an appropriate device.
         return images_adv.to(self.device)
 
+
 class FGSMNative:
     """
     Class for manually implemented FGSM, unlike the above FGSM class in this
@@ -65,9 +70,10 @@ def __init__(self, model, loss_criterion, norm=np.inf, batch_size=128):
         self.loss_criterion = loss_criterion
         self.norm = norm
         self.batch_size = batch_size
-        
+
         # Use GPU for computation if it is available
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu")
 
     def generatePerturbation(self, data, budget, minimal=False):
         """
@@ -91,7 +97,7 @@ def generatePerturbation(self, data, budget, minimal=False):
 
         images, labels = data
         images_adv = images.clone().detach().to(self.device)
-        # We will never need to compute a gradient with respect to images_adv. 
+        # We will never need to compute a gradient with respect to images_adv.
         images_adv.requires_grad_(False)
 
         images.requires_grad_(True)
@@ -104,31 +110,35 @@ def generatePerturbation(self, data, budget, minimal=False):
             direction = images.grad.data.sign()
         elif self.norm == 2:
             flattened_images = images_adv.view(-1, img_rows * img_cols)
-            direction = F.normalize(flattened_images, p=2, dim=1).view(images.size())
+            direction = F.normalize(
+                flattened_images, p=2, dim=1).view(images.size())
         else:
             raise ValueError("The norm is not valid.")
-        
+
         if minimal:
             iterations = 50
             incremental_size = budget / iterations
             minimal_perturbations = torch.zeros(images.size())
             for i in range(iterations):
-                outputs = self.model((images_adv + minimal_perturbations).clamp(0, 1))
+                outputs = self.model(
+                    (images_adv + minimal_perturbations).clamp(0, 1))
                 _, predicted = torch.max(outputs.data, 1)
                 for j in range(labels.size()[0]):
                     # If the current adversarial exampels are correctly
-                    # classified, increase the size of the perturbations. 
+                    # classified, increase the size of the perturbations.
                     if predicted[j] == labels[j]:
-                        minimal_perturbations[j].add_(incremental_size * direction[j])
+                        minimal_perturbations[j].add_(
+                            incremental_size * direction[j])
             images_adv.add_(minimal_perturbations)
         else:
             images_adv.add_(budget * direction)
-        
-        images_adv.clamp_(0,1)
 
-        # The output to be returned should be loaded on an appropriate device. 
+        images_adv.clamp_(0, 1)
+
+        # The output to be returned should be loaded on an appropriate device.
         return images_adv
 
+
 class PGD:
     """
     Module for adversarial attacks based on projected gradient descent (PGD).
@@ -144,30 +154,34 @@ def __init__(self, model, loss_criterion, norm=np.inf, batch_size=128):
         self.pytorch_model = wrapModel(model, loss_criterion)
         self.norm = norm
         self.batch_size = batch_size
-        self.attack = ProjectedGradientDescent(self.pytorch_model, norm=norm, random_init=False, batch_size=batch_size)
+        self.attack = ProjectedGradientDescent(
+            self.pytorch_model, norm=norm, random_init=False, batch_size=batch_size)
 
         # Use GPU for computation if it is available
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu")
 
     def generatePerturbation(self, data, budget, max_iter=15):
         images, _ = data
-        
-        # eps_step is not allowed to be larger than budget according to the 
-        # documentation of ART. 
+
+        # eps_step is not allowed to be larger than budget according to the
+        # documentation of ART.
         eps_step = budget / 5
-        images_adv = self.attack.generate(x=images.cpu().numpy(), norm=self.norm, eps=budget, eps_step=eps_step, max_iter=max_iter, batch_size=self.batch_size)
+        images_adv = self.attack.generate(x=images.cpu().numpy(
+        ), norm=self.norm, eps=budget, eps_step=eps_step, max_iter=max_iter, batch_size=self.batch_size)
         images_adv = torch.from_numpy(images_adv)
 
-        # The output to be returned should be loaded on an appropriate device. 
+        # The output to be returned should be loaded on an appropriate device.
         return images_adv.to(self.device)
 
+
 if __name__ == "__main__":
     # Load a simple neural network
     model = SimpleNeuralNet()
     loadModel(model, "./ERM_models/SimpleModel.pt")
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model.to(device) # Load the neural network on GPU if it is available
+    model.to(device)  # Load the neural network on GPU if it is available
     print("The neural network is now loaded on {}.".format(device))
 
     # Create an object for PGD
@@ -189,7 +203,7 @@ def generatePerturbation(self, data, budget, max_iter=15):
         # images_adv is already loaded on GPU by generatePerturbation
         images_adv = pgd.generatePerturbation(data, epsilon)
         with torch.no_grad():
-            outputs =  model(images_adv)
+            outputs = model(images_adv)
         _, predicted = torch.max(outputs.data, 1)
         total += labels.size(0)
         correct += (predicted == labels).sum().item()

adversarial_attack_DRO.py

@@ -7,6 +7,7 @@
 from util_model import SimpleNeuralNet, MNISTClassifier
 from adversarial_training import AdversarialTraining, ProjectedGradientTraining
 
+
 class ProjetcedDRO(AdversarialTraining):
     """
     Execute distributionally robust optimization (DRO) using the Euclidean
@@ -25,13 +26,13 @@ def attack(self, budget, data, steps=15):
         images_adv = images.clone().detach().to(self.device)
         images_adv.requires_grad_(True)
 
-        # images.size()[0] corresponds to the batch size. 
+        # images.size()[0] corresponds to the batch size.
         desirable_distance = budget * math.sqrt(images.size()[0])
 
         # Choose a random strating point where the constraint for perturbations
         # is tight. Without randomly choosing a starting point, the adversarial
         # attack fails most of the time because the loss function is flat near
-        # the training input, which was used in training the neural network.  
+        # the training input, which was used in training the neural network.
         randomStart(images_adv, budget)
         for i in range(steps):
             if images_adv.grad is not None:
@@ -45,13 +46,15 @@ def attack(self, budget, data, steps=15):
             distance = torch.norm(diff_tensor, p=2).item()
 
             # Inside this conditional statement, we can be certain that
-            # distance > 0, provided that budget > 0. 
-            # Hence, there is no risk of division by 0. 
+            # distance > 0, provided that budget > 0.
+            # Hence, there is no risk of division by 0.
             if distance > desirable_distance:
-                images_adv.data.add_((1 - (desirable_distance / distance)) * diff_tensor)
-            images_adv.data.clamp_(0, 1)     
+                images_adv.data.add_(
+                    (1 - (desirable_distance / distance)) * diff_tensor)
+            images_adv.data.clamp_(0, 1)
         return images_adv, labels
 
+
 class LagrangianDRO(AdversarialTraining):
     """
     Execute DRO using the Lagrangian relaxation of the original theoretical
@@ -78,21 +81,23 @@ def attack(self, budget, data, steps=15):
             budget: gamma in the original paper. Note that this parameter is
                 different from the budget parameter in other DRO classes. 
         """
-        
+
         images, labels = data
         images_adv = images.clone().detach().to(self.device)
         images_adv.requires_grad_(True)
-        
+
         for i in range(steps):
             if images_adv.grad is not None:
                 images_adv.grad.data.zero_()
             outputs = self.model(images_adv)
-            loss = self.loss_criterion(outputs, labels) - budget * self.cost_function(images, images_adv)
+            loss = self.loss_criterion(
+                outputs, labels) - budget * self.cost_function(images, images_adv)
             loss.backward()
             images_adv.data.add_(1 / math.sqrt(i+1) * images_adv.grad)
             images_adv.data.clamp_(0, 1)
         return images_adv, labels
 
+
 class FrankWolfeDRO(AdversarialTraining):
     """
     Execute DRO using the Frank-Wolfe method together with the stochastic
@@ -124,19 +129,21 @@ def attack(self, budget, data, steps=15):
         images, labels = data
         images_adv = images.clone().detach().to(self.device)
         images_adv.requires_grad_(True)
-        
+
         for i in range(steps):
             if images_adv.grad is not None:
                 images_adv.grad.zero_()
             outputs = self.model(images_adv)
             loss = self.loss_criterion(outputs, labels)
             loss.backward()
 
-            # desitnation corresponds to y_t in the paper by Bubeck. 
-            destination = images_adv.data + self.getOptimalDirection(budget=budget, data=images_adv.grad)
+            # desitnation corresponds to y_t in the paper by Bubeck.
+            destination = images_adv.data + \
+                self.getOptimalDirection(budget=budget, data=images_adv.grad)
             destination = destination.to(self.device)
             gamma = 2 / (i + 2)
-            images_adv.data = (1 - gamma) * images_adv.data + gamma * destination
+            images_adv.data = (1 - gamma) * \
+                images_adv.data + gamma * destination
             images_adv.data.clamp_(0, 1)
         return images_adv, labels
 
@@ -154,15 +161,15 @@ def getOptimalDirection(self, budget, data):
             data: gradient of the total loss with respect to the current
                 batch of adversarial examples. This corresponds to C in
                 Appendix B of the paper by Staib et al. 
-        
+
         Returns:
             X in Appendix B of Staib et al.'s paper 
         """
 
         # The number of samples
         batch_size = data.size()[0]
 
-        # 'directions' corresponds to v's in Staib et al.'s paper. 
+        # 'directions' corresponds to v's in Staib et al.'s paper.
         directions = data.clone().detach().view((batch_size, -1))
         directions = directions.to(self.device)
 
@@ -173,17 +180,17 @@ def getOptimalDirection(self, budget, data):
             directions.pow_(normalize_dim)
             directions = F.normalize(directions, p=self.q, dim=1)
         else:
-            raise ValueError("The value of q must be larger than 1.")      
-        
-        # This corresponds to a's in the original paper. 
+            raise ValueError("The value of q must be larger than 1.")
+
+        # This corresponds to a's in the original paper.
         products = []
         for i, direction in enumerate(directions):
             sample = data[i].view(-1)
             products.append(torch.dot(direction, sample))
         products = torch.stack(products)
         products = products.to(self.device)
 
-        # This corresponds to epsilons in the original paper. 
+        # This corresponds to epsilons in the original paper.
         size_factors = products.clone().detach()
         size_factors = size_factors.to(self.device)
         if self.p == np.inf:
@@ -192,16 +199,17 @@ def getOptimalDirection(self, budget, data):
             normalize_dim = 1 / (self.p - 1)
             size_factors.pow_(normalize_dim)
             distance = torch.norm(size_factors, p=self.p).item()
-            size_factors = size_factors / distance # This is now normalized. 
+            size_factors = size_factors / distance  # This is now normalized.
         else:
-            raise ValueError("The value of p must be larger than 1.")  
-        
+            raise ValueError("The value of p must be larger than 1.")
+
         outputs = []
         for i, size_factor in enumerate(size_factors):
             outputs.append(directions[i] * size_factor * budget)
         result = torch.stack(outputs).view(data.size())
         return result.to(self.device)
 
+
 def trainDROModel(dro_type, epochs, steps_adv, budget, activation, batch_size, loss_criterion, cost_function=None):
     """
     Train a neural network using one of the following DRO methods:
@@ -210,7 +218,7 @@ def trainDROModel(dro_type, epochs, steps_adv, budget, activation, batch_size, l
             This is also called WRM.
         - the Frank-Wolfe method based approach developed by Staib et al. 
     """
-    
+
     model = MNISTClassifier(activation=activation)
     if dro_type == 'PGD':
         train_module = ProjetcedDRO(model, loss_criterion)
@@ -222,11 +230,16 @@ def trainDROModel(dro_type, epochs, steps_adv, budget, activation, batch_size, l
     else:
         raise ValueError("The type of DRO is not valid.")
 
-    train_module.train(budget=budget, batch_size=batch_size, epochs=epochs, steps_adv=steps_adv)
+    train_module.train(budget=budget, batch_size=batch_size,
+                       epochs=epochs, steps_adv=steps_adv)
     folderpath = "./DRO_models/"
-    filepath = folderpath + "{}_DRO_activation={}_epsilon={}.pt".format(dro_type, activation, budget)
+    filepath = folderpath + \
+        "{}_DRO_activation={}_epsilon={}.pt".format(
+            dro_type, activation, budget)
     torch.save(model.state_dict(), filepath)
-    print("A neural network adversarially trained using {} is now saved at {}.".format(dro_type, filepath))
+    print("A neural network adversarially trained using {} is now saved at {}.".format(
+        dro_type, filepath))
+
 
 if __name__ == "__main__":
     epochs = 25
@@ -235,14 +248,21 @@ def trainDROModel(dro_type, epochs, steps_adv, budget, activation, batch_size, l
     gammas = [0.0001, 0.0003, 0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1.0, 3.0]
     batch_size = 128
     loss_criterion = nn.CrossEntropyLoss()
-    cost_function = lambda x, y: torch.dist(x, y, p=2) ** 2
 
-    trainDROModel('PGD', epochs, steps_adv, epsilon, 'relu', batch_size, loss_criterion)
-    trainDROModel('FW', epochs, steps_adv, epsilon, 'relu', batch_size, loss_criterion)
+    def cost_function(x, y): return torch.dist(x, y, p=2) ** 2
+
+    trainDROModel('PGD', epochs, steps_adv, epsilon,
+                  'relu', batch_size, loss_criterion)
+    trainDROModel('FW', epochs, steps_adv, epsilon,
+                  'relu', batch_size, loss_criterion)
 
-    trainDROModel('PGD', epochs, steps_adv, epsilon, 'elu', batch_size, loss_criterion)
-    trainDROModel('FW', epochs, steps_adv, epsilon, 'elu', batch_size, loss_criterion)
+    trainDROModel('PGD', epochs, steps_adv, epsilon,
+                  'elu', batch_size, loss_criterion)
+    trainDROModel('FW', epochs, steps_adv, epsilon,
+                  'elu', batch_size, loss_criterion)
 
     for gamma in gammas:
-        trainDROModel('Lag', epochs, steps_adv, gamma, 'relu', batch_size, loss_criterion, cost_function=cost_function)
-        trainDROModel('Lag', epochs, steps_adv, gamma, 'elu', batch_size, loss_criterion, cost_function=cost_function)
+        trainDROModel('Lag', epochs, steps_adv, gamma, 'relu',
+                      batch_size, loss_criterion, cost_function=cost_function)
+        trainDROModel('Lag', epochs, steps_adv, gamma, 'elu',
+                      batch_size, loss_criterion, cost_function=cost_function)