add adversarial taining

README.md

@@ -1,21 +1,23 @@
+Implementation of some unbalanced loss for NLP task like focal_loss, dice_loss, DSC Loss, GHM Loss et.al and adversarial training like FGM, FGSM, PGD, FreeAT.  
 
+### Loss Summary
 
-Implementation of some unbalanced loss for NLP task like focal_loss, dice_loss, DSC Loss, GHM Loss et.al
-### Summary
 Here is a loss implementation repository included unbalanced loss
 
-| Loss Name | paper | Notes |
-| :-----:| :----: | :----: |
-| Weighted CE Loss | [UNet Architectures in Multiplanar Volumetric Segmentation -- Validated on Three Knee MRI Cohorts](https://arxiv.org/abs/1708.02002) |  |
-| Focal Loss | [Focal Loss for Dense Object Detection](https://arxiv.org/abs/2203.08194) |  |
-| Dice Loss | [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://arxiv.org/abs/1606.04797) |  |
-| DSC Loss | [Dice Loss for Data-imbalanced NLP Tasks](https://arxiv.org/pdf/1911.02855.pdf) |  |
-| GHM Loss | [Gradient Harmonized Single-stage Detector](https://www.aaai.org/ojs/index.php/AAAI/article/download/4877/4750) |  |
-
-### How to use?
-You can find all the loss usage information in test_loss.py.  
-
+|    Loss Name     |                                                                paper                                                                 | Notes |
+|:----------------:|:------------------------------------------------------------------------------------------------------------------------------------:|:-----:|
+| Weighted CE Loss | [UNet Architectures in Multiplanar Volumetric Segmentation -- Validated on Three Knee MRI Cohorts](https://arxiv.org/abs/2203.08194) |       |
+|    Focal Loss    |                              [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002)                               |       |
+|    Dice Loss     |       [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://arxiv.org/abs/1606.04797)       |       |
+|     DSC Loss     |                           [Dice Loss for Data-imbalanced NLP Tasks](https://arxiv.org/pdf/1911.02855.pdf)                            |       |
+|     GHM Loss     |           [Gradient Harmonized Single-stage Detector](https://www.aaai.org/ojs/index.php/AAAI/article/download/4877/4750)            |       |
+
+#### How to use?
+
+You can find all the loss usage information in test_loss.py.
+
 Here is a simple demo of usage:
+
 ```python
 import torch
 from unbalanced_loss.focal_loss import MultiFocalLoss
@@ -24,8 +26,51 @@ batch_size, num_class = 64, 10
 Loss_Func = MultiFocalLoss(num_class=num_class, gamma=2.0, reduction='mean')
 
 logits = torch.rand(batch_size, num_class, requires_grad=True)  # (batch_size, num_classes)
-targets = torch.randint(0, num_class, size=(batch_size, ))  # (batch_size, )
+targets = torch.randint(0, num_class, size=(batch_size,))  # (batch_size, )
 
 loss = Loss_Func(logits, targets)
 loss.backward()
-```
+```
+
+### Adversarial Training Summary
+
+Here is a Summary of Adversarial Training implementation.   
+you can find more details in adversarial_training/README.md
+
+| Adversarial Training |                                                 paper                                                 | Notes |
+|:--------------------:|:-----------------------------------------------------------------------------------------------------:|:-----:|
+|         FGM          |                     [Fast Gradient Method](https://arxiv.org/pdf/1605.07725.pdf)                      |       |
+|         FGSM         |                     [Fast Gradient Sign Method](https://arxiv.org/abs/1412.6572)                      |       |
+|         PGD          | [Towards Deep Learning Models Resistant to Adversarial Attacks](https://arxiv.org/pdf/1706.06083.pdf) |       |
+|        FreeAT        |                   [Free Adversarial Training](https://arxiv.org/pdf/1904.12843.pdf)                   |       |
+|        FreeLB        |            [Free Large Batch Adversarial Training](https://arxiv.org/pdf/1909.11764v5.pdf)            |       |
+
+#### How to use?  
+You can find a simple demo for bert classification in test_bert.py.  
+
+Here is a simple demo of usage:  
+You just need to rewrite train function according to input for your model in file PGD.py, then you can use adversarial training like below.  
+```python
+import transformers
+from model import bert_classification
+from adversarial_training.PGD import PGD
+
+batch_size, num_class = 64, 10
+# model = your_model()
+model = bert_classification()
+AT_Model = PGD(model)
+optimizer = transformers.AdamW(model.parameters(), lr=0.001)
+
+# rewrite your train function in pgd.py
+outputs, loss = AT_Model.train_bert(token, segment, mask, label, optimizer)
+```
+
+#### Adversarial Training Results Compare
+
+|  Adversarial Training  | Time Cost(s/epoch ) | best_acc |
+|:----------------------:|:-------------------:|:--------:|
+| Normal(not add attack) |        23.77        |  0.773   |
+|          FGSM          |        45.95        |  0.7936  |
+|          FGM           |        47.28        |  0.8008  |
+|        PGD(k=3)        |        87.50        |  0.7963  |
+|      FreeAT(k=3)       |        93.26        |  0.7896  |

adversarial_training/FGM.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn.functional as F
+
+
+class FGM:
+    def __init__(self, model, emb_name='embedding.'):
+        self.model = model
+        self.emb_backup = {}  # restore embedding parameters
+        self.epsilon = 1.0
+        self.emb_name = emb_name
+
+    def train(self, input_data, labels, optimizer):
+        ''' define process of training here according to your model define
+        '''
+        pass
+
+    def train_bert(self, token, segment, mask, labels, optimizer, attack=False):
+        ''' a advertisement training demo for bert
+        '''
+        outputs = self.model(token, segment, mask)
+        loss = F.cross_entropy(outputs, labels)
+        loss.backward()
+
+        if attack:
+            self.attack_embedding()
+            outputs = self.model(token, segment, mask)
+            loss = F.cross_entropy(outputs, labels)
+            # self.model.zero_grad()  # compute advertise samples' grad only
+            loss.backward()
+            self.restore_embedding()  # recover
+        optimizer.step()
+        self.model.zero_grad()
+
+        return outputs, loss
+
+    def attack_embedding(self, backup=True):
+        ''' add add disturbance in embedding layer you want
+        '''
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                if backup:  # store parameter
+                    self.emb_backup[name] = param.data.clone()
+
+                self._add_disturbance(name, param)  # add disturbance
+
+    def restore_embedding(self):
+        '''recover embedding backup before
+        '''
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                assert name in self.emb_backup
+                param.data = self.emb_backup[name]
+        self.emb_backup = {}
+
+    def _add_disturbance(self, name, param):
+        ''' add disturbance
+        '''
+        norm = torch.norm(param.grad)
+        if norm != 0:
+            r_at = self.epsilon * param.grad / norm
+            param.data.add_(r_at)

adversarial_training/FGSM.py

@@ -0,0 +1,54 @@
+import numpy as np
+import torch.nn.functional as F
+
+
+class FGSM:
+    def __init__(self, model, epsilon=0.05, emb_name='embedding.'):
+        self.model = model
+        self.emb_backup = {}
+        self.epsilon = epsilon
+        self.emb_name = emb_name
+
+    def train(self, input_data, labels, optimizer):
+        ''' define process of training here according to your model define
+        '''
+        pass
+
+    def train_bert(self, token, segment, mask, labels, optimizer, attack=False):
+        ''' add disturbance in training
+        '''
+        outputs = self.model(token, segment, mask)
+        loss = F.cross_entropy(outputs, labels)
+        loss.backward()
+
+        if attack:
+            self.attack_embedding()
+            outputs = self.model(token, segment, mask)
+            loss = F.cross_entropy(outputs, labels)
+            # self.model.zero_grad()  # compute advertise samples' grad only
+            loss.backward()
+            self.restore_embedding()  # recover
+        optimizer.step()
+        self.model.zero_grad()
+
+        return outputs, loss
+
+    def attack_param(self, name, param):
+        # r_at = epsilon * sign(grad)
+        r_at = self.epsilon * np.sign(param.grad)
+        param.data.add_(r_at)
+
+    def attack_embedding(self, backup=True):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                if backup:
+                    self.emb_backup[name] = param.data.clone()
+                # attack embedding
+                self.attack_param(name, param)
+
+    def restore_embedding(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                assert name in self.emb_backup
+                param.data = self.emb_backup[name]
+        self.emb_backup = {}

adversarial_training/FreeAT.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn.functional as F
+
+
+
+class FreeAT:
+    def __init__(self, model, epsilon=0.8, k=3, emb_name='embedding.'):
+        self.model = model
+        self.emb_backup = {}
+        self.epsilon = epsilon
+        self.K = k  # attack times
+        self.emb_name = emb_name  # embedding layer name want to attack
+        self.backup_emb()
+
+    def train(self, input_data, labels, optimizer):
+        ''' define process of training here according to your model define
+        '''
+        pass
+
+    def train_bert(self, token, segment, mask, labels, optimizer, attack=True):
+        ''' add disturbance in training
+        '''
+        outputs = self.model(token, segment, mask)
+        loss = F.cross_entropy(outputs, labels)
+        loss.backward()
+
+        if attack:
+            for t in range(self.K):
+                outputs = self.model(token, segment, mask)
+                self.model.zero_grad()
+                loss = F.cross_entropy(outputs, labels)
+                loss.backward()
+                optimizer.step()
+                self.attack_emb(backup=False)  # accumulate projected disturb in embedding
+
+        return outputs, loss
+
+    def attack_param(self, name, param):
+        '''add  disturbance
+           FreeAT Format:
+            r[t+1] = r[t] + epsilon * sign(grad)
+            r_at = epsilon * np.sign(param.grad)
+        '''
+        norm = torch.norm(param.grad)
+        if norm != 0:
+            r_at = self.epsilon * param.grad / norm
+            param.data.add_(r_at)
+            param.data = self.project(name, param.data)
+
+    def project(self, param_name, param_data):
+        ''' projected  disturbance like disturb cropping inside the pale (-eps, eps)
+        '''
+        r = param_data - self.emb_backup[param_name]  # compute  disturbance
+        if torch.norm(r) > self.epsilon:  # disturbance cropping inside the pale (-eps, eps)
+            r = self.epsilon * r / torch.norm(r)
+        return self.emb_backup[param_name] + r
+
+    def attack_emb(self, backup=False):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                if backup:  # backup embedding
+                    self.emb_backup[name] = param.data.clone()
+                self.attack_param(name, param)
+
+    def backup_emb(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                self.emb_backup[name] = param.data.clone()
+
+    def restore_emb(self):
+        '''recover embedding'''
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                assert name in self.emb_backup
+                param.data = self.emb_backup[name]
+        self.emb_backup = {}

adversarial_training/PGD.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn.functional as F
+
+
+class PGD:
+    def __init__(self, model, epsilon=1.0, alpha=0.3, k=3, emb_name='embedding.'):
+        self.model = model
+        self.emb_backup = {}
+        self.grad_backup = {}
+        self.epsilon = epsilon
+        self.alpha = alpha
+        self.K = k  # PGD attack times
+        self.emb_name = emb_name
+
+    def train(self, input_data, labels, optimizer):
+        ''' define process of training here according to your model define
+        '''
+        pass
+
+    def train_bert(self, token, segment, mask, labels, optimizer, attack=True):
+        ''' a advertisement training demo for bert
+        '''
+        outputs = self.model(token, segment, mask)
+        loss = F.cross_entropy(outputs, labels)
+        loss.backward()
+
+        if attack:
+            self.backup_grad()
+            for t in range(self.K):
+                self.attack_embedding(backup=(t == 0))
+                if t != self.K - 1:
+                    self.model.zero_grad()
+                else:
+                    self.restore_grad()
+                outputs = self.model(token, segment, mask)
+                loss = F.cross_entropy(outputs, labels)
+                loss.backward()
+            self.restore_embedding()  # recover embedding
+        optimizer.step()
+        self.model.zero_grad()
+
+        return outputs, loss
+
+    def attack_param(self, name, param):
+        '''add  disturbance
+            PGD: r = epsilon * grad / norm(grad)
+        '''
+        norm = torch.norm(param.grad)
+        if norm != 0 and not torch.isnan(norm):
+            r_at = self.alpha * param.grad / norm
+            param.data.add_(r_at)
+            param.data = self.project(name, param.data)
+
+    def project(self, param_name, param_data):
+        ''' projected  disturbance like parameter cropping inside the pale
+        '''
+        r = param_data - self.emb_backup[param_name]
+        if torch.norm(r) > self.epsilon:
+            r = self.epsilon * r / torch.norm(r)
+        return self.emb_backup[param_name] + r
+
+    def attack_embedding(self, backup=False):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                if backup:  # backup embedding
+                    self.emb_backup[name] = param.data.clone()
+                self.attack_param(name, param)
+
+    def backup_embedding(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                self.emb_backup[name] = param.data.clone()
+
+    def restore_embedding(self):
+        '''recover embedding'''
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and self.emb_name in name:
+                assert name in self.emb_backup
+                param.data = self.emb_backup[name]
+        self.emb_backup = {}
+
+    def backup_grad(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and param.grad is not None:
+                self.grad_backup[name] = param.grad.clone()
+
+    def restore_grad(self):
+        '''recover grad back upped
+        '''
+        for name, param in self.model.named_parameters():
+            if param.requires_grad and name in self.grad_backup:
+                param.grad = self.grad_backup[name]