update

Author: namisan

mt_dnn/batcher.py

@@ -257,7 +257,7 @@ def __init__(self,
                  max_seq_length=512,
                  max_predictions_per_seq=80,
                  printable=True):
-        data, tokenizer = self.load(path, is_train, maxlen, factor, task_def, bert_model, do_lower_case)
+        data, tokenizer = self.load(path, is_train, maxlen, factor, task_def, bert_model, do_lower_case, printable=printable)
         self._data = data
         self._tokenizer = tokenizer
         self._task_id = task_id

mt_dnn/loss.py

@@ -4,6 +4,7 @@
 import torch
 from torch.nn.modules.loss import _Loss
 import torch.nn.functional as F
+import torch.nn as nn
 from enum import IntEnum
 
 def stable_kl(logit, target, epsilon=1e-6, reduce=True):
@@ -49,6 +50,7 @@ def forward(self, input, target, weight=None, ignore_index=-1):
         loss = loss * self.alpha
         return loss
 
+
 class SeqCeCriterion(CeCriterion):
     def __init__(self, alpha=1.0, name='Seq Cross Entropy Criterion'):
         super().__init__(alpha, name)
@@ -116,13 +118,13 @@ def __init__(self, alpha=1.0, name='KL Div Criterion'):
         self.alpha = alpha
         self.name = name
 
-    def forward(self, input, target, weight=None, ignore_index=-1):
+    def forward(self, input, target, weight=None, ignore_index=-1, reduction='batchmean'):
         """input/target: logits
         """
         input = input.float()
         target = target.float()
-        loss = F.kl_div(F.log_softmax(input, dim=-1, dtype=torch.float32), F.softmax(target.detach(), dim=-1, dtype=torch.float32), reduction='batchmean') + \
-            F.kl_div(F.log_softmax(target, dim=-1, dtype=torch.float32), F.softmax(input.detach(), dim=-1, dtype=torch.float32), reduction='batchmean')
+        loss = F.kl_div(F.log_softmax(input, dim=-1, dtype=torch.float32), F.softmax(target.detach(), dim=-1, dtype=torch.float32), reduction=reduction) + \
+            F.kl_div(F.log_softmax(target, dim=-1, dtype=torch.float32), F.softmax(input.detach(), dim=-1, dtype=torch.float32), reduction=reduction)
         loss = loss * self.alpha
         return loss
 
@@ -142,6 +144,41 @@ def forward(self, input, target, weight=None, ignore_index=-1):
         loss = loss * self.alpha
         return loss
 
+class JSCriterion(Criterion):
+    def __init__(self, alpha=1.0, name='JS Div Criterion'):
+        super().__init__()
+        self.alpha = alpha
+        self.name = name
+
+    def forward(self, input, target, weight=None, ignore_index=-1, reduction='batchmean'):
+        """input/target: logits
+        """
+        input = input.float()
+        target = target.float()
+        m = F.softmax(target.detach(), dim=-1, dtype=torch.float32) + \
+            F.softmax(input.detach(), dim=-1, dtype=torch.float32)
+        m = 0.5 * m
+        loss = F.kl_div(F.log_softmax(input, dim=-1, dtype=torch.float32), m, reduction=reduction) + \
+            F.kl_div(F.log_softmax(target, dim=-1, dtype=torch.float32), m, reduction=reduction)
+        loss = loss * self.alpha
+        return loss
+
+class HLCriterion(Criterion):
+    def __init__(self, alpha=1.0, name='Hellinger Criterion'):
+        super().__init__()
+        self.alpha = alpha
+        self.name = name
+
+    def forward(self, input, target, weight=None, ignore_index=-1, reduction='batchmean'):
+        """input/target: logits
+        """
+        input = input.float()
+        target = target.float()
+        si = F.softmax(target.detach(), dim=-1, dtype=torch.float32).sqrt_()
+        st = F.softmax(input.detach(), dim=-1, dtype=torch.float32).sqrt_()
+        loss = F.mse_loss(si, st)
+        loss = loss * self.alpha
+        return loss
 
 
 class RankCeCriterion(Criterion):
@@ -202,42 +239,6 @@ def forward(self, input, target, weight=None, ignore_index=-1):
         loss = loss * self.alpha
         return loss
 
-class JSCriterion(Criterion):
-    def __init__(self, alpha=1.0, name='JS Div Criterion'):
-        super().__init__()
-        self.alpha = alpha
-        self.name = name
-
-    def forward(self, input, target, weight=None, ignore_index=-1, reduction='batchmean'):
-        """input/target: logits
-        """
-        input = input.float()
-        target = target.float()
-        m = F.softmax(target.detach(), dim=-1, dtype=torch.float32) + \
-            F.softmax(input.detach(), dim=-1, dtype=torch.float32)
-        m = 0.5 * m
-        loss = F.kl_div(F.log_softmax(input, dim=-1, dtype=torch.float32), m, reduction=reduction) + \
-            F.kl_div(F.log_softmax(target, dim=-1, dtype=torch.float32), m, reduction=reduction)
-        loss = loss * self.alpha
-        return loss
-
-class HLCriterion(Criterion):
-    def __init__(self, alpha=1.0, name='Hellinger Criterion'):
-        super().__init__()
-        self.alpha = alpha
-        self.name = name
-
-    def forward(self, input, target, weight=None, ignore_index=-1, reduction='batchmean'):
-        """input/target: logits
-        """
-        input = input.float()
-        target = target.float()
-        si = F.softmax(target.detach(), dim=-1, dtype=torch.float32).sqrt_()
-        st = F.softmax(input.detach(), dim=-1, dtype=torch.float32).sqrt_()
-        loss = F.mse_loss(si, st)
-        loss = loss * self.alpha
-        return loss
-
 class LossCriterion(IntEnum):
     CeCriterion = 0
     MseCriterion = 1
@@ -252,6 +253,7 @@ class LossCriterion(IntEnum):
     JSCriterion = 10
     HLCriterion = 11
 
+
 LOSS_REGISTRY = {
      LossCriterion.CeCriterion: CeCriterion,
      LossCriterion.MseCriterion: MseCriterion,

mt_dnn/perturbation.py

@@ -88,7 +88,7 @@ def forward(self, model,
                 if task_type == TaskType.Ranking:
                     adv_logits = adv_logits.view(-1, pairwise)
                 adv_loss = stable_kl(adv_logits, logits.detach(), reduce=False) 
-            delta_grad, = torch.autograd.grad(adv_loss, noise, only_inputs=True)
+            delta_grad, = torch.autograd.grad(adv_loss, noise, only_inputs=True, retain_graph=False)
             norm = delta_grad.norm()
             if (torch.isnan(norm) or torch.isinf(norm)):
                 return 0

train.py

@@ -145,7 +145,7 @@ def train_config(parser):
     parser.add_argument('--adv_p_norm', default='inf', type=str)
     parser.add_argument('--adv_alpha', default=1, type=float)
     parser.add_argument('--adv_k', default=1, type=int)
-    parser.add_argument('--adv_step_size', default=1e-3, type=float)
+    parser.add_argument('--adv_step_size', default=1e-5, type=float)
     parser.add_argument('--adv_noise_var', default=1e-5, type=float)
     parser.add_argument('--adv_epsilon', default=1e-6, type=float)
     parser.add_argument('--encode_mode', action='store_true', help="only encode test data")