Upload folder using huggingface_hub

robust_loss/__init__.py

@@ -0,0 +1,7 @@
+from .hat import *
+from .mart import *
+from .trades import *
+from .rslad import *
+from .sat import *
+from .ard import *
+from .adaad import *

robust_loss/adaad.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import torch.optim as optim
+import numpy as np
+def adaad_inner_loss(model,
+                     teacher_model,
+                     x_natural,
+                     step_size=2/255,
+                     steps=10,
+                     epsilon=8/255,
+                     BN_eval=True,
+                     random_init=True,
+                     clip_min=0.0,
+                     clip_max=1.0):
+    # define KL-loss
+    criterion_kl = nn.KLDivLoss(reduction='none')
+    if BN_eval:
+        model.eval()
+
+    # set eval mode for teacher model
+    teacher_model.eval()
+    # generate adversarial example
+    if random_init:
+        x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cuda().detach()
+    else:
+        x_adv = x_natural.detach()
+    for _ in range(steps):
+        x_adv.requires_grad_()
+        with torch.enable_grad():
+            loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
+                                   F.softmax(teacher_model(x_adv), dim=1))
+            loss_kl = torch.sum(loss_kl)
+        grad = torch.autograd.grad(loss_kl, [x_adv])[0]
+        x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
+        x_adv = torch.min(torch.max(x_adv, x_natural -
+                          epsilon), x_natural + epsilon)
+        x_adv = torch.clamp(x_adv, clip_min, clip_max)
+
+    if BN_eval:
+        model.train()
+    model.train()
+
+    x_adv = Variable(torch.clamp(x_adv, clip_min, clip_max),
+                     requires_grad=False)
+    return x_adv
+def adaad_loss(teacher_model,model,x_natural,y,optimizer,step_size=0.0078,
+                epsilon=0.031,
+                perturb_steps=10,
+                beta = 6.0,
+                AdaAD_alpha=1.0):
+    adv_inputs  = adaad_inner_loss(model,teacher_model,x_natural,step_size,perturb_steps,epsilon)
+    ori_outputs = model(x_natural)
+    adv_outputs = model(adv_inputs)
+    with torch.no_grad():
+        teacher_model.eval()
+        t_ori_outputs = teacher_model(x_natural)
+        t_adv_outputs = teacher_model(adv_inputs)
+    kl_loss1 = nn.KLDivLoss()(F.log_softmax(adv_outputs, dim=1),
+                                          F.softmax(t_adv_outputs.detach(), dim=1))
+    kl_loss2 = nn.KLDivLoss()(F.log_softmax(ori_outputs, dim=1),
+                                          F.softmax(t_ori_outputs.detach(), dim=1))
+    loss = AdaAD_alpha*kl_loss1 + (1-AdaAD_alpha)*kl_loss2
+    return loss

robust_loss/ard.py

@@ -0,0 +1,72 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import torch.optim as optim
+import numpy as np
+def attack_pgd(model,train_batch_data,train_batch_labels,attack_iters=10,step_size=2/255.0,epsilon=8.0/255.0):
+    ce_loss = torch.nn.CrossEntropyLoss().cuda()
+    train_ifgsm_data = train_batch_data.detach() + torch.zeros_like(train_batch_data).uniform_(-epsilon,epsilon)
+    train_ifgsm_data = torch.clamp(train_ifgsm_data,0,1)
+    for i in range(attack_iters):
+        train_ifgsm_data.requires_grad_()
+        logits = model(train_ifgsm_data)
+        loss = ce_loss(logits,train_batch_labels.cuda())
+        loss.backward()
+        train_grad = train_ifgsm_data.grad.detach()
+        train_ifgsm_data = train_ifgsm_data + step_size*torch.sign(train_grad)
+        train_ifgsm_data = torch.clamp(train_ifgsm_data.detach(),0,1)
+        train_ifgsm_pert = train_ifgsm_data - train_batch_data
+        train_ifgsm_pert = torch.clamp(train_ifgsm_pert,-epsilon,epsilon)
+        train_ifgsm_data = train_batch_data + train_ifgsm_pert
+        train_ifgsm_data = train_ifgsm_data.detach()
+    return train_ifgsm_data
+
+def ard_inner_loss(model,
+                teacher_logits,
+                x_natural,
+                y,
+                optimizer,
+                step_size=0.0078,
+                epsilon=0.031,
+                perturb_steps=10,
+                beta=6.0):
+    # define KL-loss
+    criterion_kl = nn.KLDivLoss(size_average=False,reduce=False)
+    model.eval()
+    batch_size = len(x_natural)
+    # generate adversarial example
+    x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cuda().detach()
+
+    for _ in range(perturb_steps):
+        x_adv.requires_grad_()
+        with torch.enable_grad():
+            loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
+                                       F.softmax(teacher_logits, dim=1))
+            loss_kl = torch.sum(loss_kl)
+        grad = torch.autograd.grad(loss_kl, [x_adv])[0]
+        x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
+        x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)
+        x_adv = torch.clamp(x_adv, 0.0, 1.0)
+
+    model.train()
+
+    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
+    # zero gradient
+    # optimizer.zero_grad()
+    logits = model(x_adv)
+    return logits
+def ard_loss(teacher_model,model,x_natural,y,optimizer,step_size=0.0078,
+                epsilon=0.031,
+                perturb_steps=10,
+                beta=6.0,
+                alpha = 1.0,
+                temp = 30.0):
+    KL_loss = nn.KLDivLoss()
+    XENT_loss = nn.CrossEntropyLoss()
+    teacher_logits = teacher_model(x_natural)
+    adv_logits = ard_inner_loss(model,teacher_logits,x_natural,y,optimizer,step_size,epsilon,perturb_steps)
+    model.train()
+    nat_logits = model(x_natural)
+    loss = alpha*temp*temp*KL_loss(F.log_softmax(adv_logits/temp, dim=1),F.softmax(teacher_logits/temp, dim=1))+(1.0-alpha)*XENT_loss(nat_logits, y)
+    return loss

robust_loss/hat.py

@@ -0,0 +1,77 @@
+import numpy as np
+from torch.autograd import Variable
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from robust_loss.sat import ctx_noparamgrad_and_eval
+def hat_loss(model, x, y, optimizer, step_size=0.007, epsilon=0.031, perturb_steps=10,  beta=1.0, 
+             attack='linf',natural_criterion= None ,h=3.5, gamma=1.0, hr_model=None):
+    """
+    TRADES + Helper-based adversarial training.
+    """
+  
+    criterion_kl = nn.KLDivLoss(reduction='sum')
+    model.eval()
+  
+    x_adv = x.detach() + 0.001 * torch.randn(x.shape).cuda().detach()
+    p_natural = F.softmax(model(x), dim=1)
+  
+    if attack == 'l_inf':
+        for _ in range(perturb_steps):
+            x_adv.requires_grad_()
+            with torch.enable_grad():
+                loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1), p_natural)
+            grad = torch.autograd.grad(loss_kl, [x_adv])[0]
+            x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
+            x_adv = torch.min(torch.max(x_adv, x - epsilon), x + epsilon)
+            x_adv = torch.clamp(x_adv, 0.0, 1.0)
+    elif attack == 'l2':
+        delta = 0.001 * torch.randn(x.shape).cuda().detach()        
+        delta = Variable(delta.data, requires_grad=True)
+        
+        batch_size = len(x)
+        optimizer_delta = torch.optim.SGD([delta], lr=step_size)
+
+        for _ in range(perturb_steps):
+            adv = x + delta
+            optimizer_delta.zero_grad()
+            with torch.enable_grad():
+                loss = (-1) * criterion_kl(F.log_softmax(model(adv), dim=1), p_natural)
+            loss.backward(retain_graph=True)
+          
+            grad_norms = delta.grad.view(batch_size, -1).norm(p=2, dim=1)
+            delta.grad.div_(grad_norms.view(-1, 1, 1, 1))
+            if (grad_norms == 0).any():
+                delta.grad[grad_norms == 0] = torch.randn_like(delta.grad[grad_norms == 0])
+            optimizer_delta.step()
+
+            delta.data.add_(x)
+            delta.data.clamp_(0, 1).sub_(x)
+            delta.data.renorm_(p=2, dim=0, maxnorm=epsilon)
+        x_adv = Variable(x + delta, requires_grad=False)
+    else:
+        raise ValueError(f'Attack={attack} not supported for TRADES training!')
+    model.train()
+
+  
+    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
+    x_hr = x + h * (x_adv - x)
+    if hr_model == None:
+        with ctx_noparamgrad_and_eval(model):
+            y_hr = model(x_adv).argmax(dim=1)
+    else:        
+        with ctx_noparamgrad_and_eval(hr_model):
+            y_hr = hr_model(x_adv).argmax(dim=1)
+  
+    optimizer.zero_grad()
+  
+    out_clean, out_adv, out_help = model(x), model(x_adv), model(x_hr)
+    loss_clean = F.cross_entropy(out_clean, y, reduction='mean')
+    loss_adv = (1/len(x)) * criterion_kl(F.log_softmax(out_adv, dim=1), F.softmax(out_clean, dim=1))
+  
+    loss_help = F.cross_entropy(out_help, y_hr, reduction='mean')
+    loss = loss_clean + beta * loss_adv + gamma * loss_help
+     
+    
+  
+    return loss

robust_loss/mart.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from utils.criterion import CrossEntropyWithLabelSmooth
+
+
+def mart_loss(model, x_natural, y, optimizer, step_size=0.007, epsilon=0.031, perturb_steps=10, beta=6.0, 
+              attack='l_inf',natural_criterion= nn.CrossEntropyLoss()):
+    """
+    MART training (Wang et al, 2020).
+    """
+  
+    kl = nn.KLDivLoss(reduction='none')
+    model.eval()
+    batch_size = len(x_natural)
+    
+    # generate adversarial example
+    x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cuda().detach()
+    if attack == 'l_inf':
+        for _ in range(perturb_steps):
+            x_adv.requires_grad_()
+            with torch.enable_grad():
+                loss_ce = natural_criterion(model(x_adv), y)
+            grad = torch.autograd.grad(loss_ce, [x_adv])[0]
+            x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
+            x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)
+            x_adv = torch.clamp(x_adv, 0.0, 1.0)
+    else:
+        raise ValueError(f'Attack={attack} not supported for MART training!')
+    model.train()
+    
+    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
+    # zero gradient
+    optimizer.zero_grad()
+
+    logits = model(x_natural)
+    logits_adv = model(x_adv)
+    
+    adv_probs = F.softmax(logits_adv, dim=1)
+    tmp1 = torch.argsort(adv_probs, dim=1)[:, -2:]
+    new_y = torch.where(tmp1[:, -1] == y, tmp1[:, -2], tmp1[:, -1])
+    loss_adv = natural_criterion(logits_adv, y) + F.nll_loss(torch.log(1.0001 - adv_probs + 1e-12), new_y)
+
+    nat_probs = F.softmax(logits, dim=1)
+    true_probs = torch.gather(nat_probs, 1, (y.unsqueeze(1)).long()).squeeze()
+
+    loss_robust = (1.0 / batch_size) * torch.sum(
+        torch.sum(kl(torch.log(adv_probs + 1e-12), nat_probs), dim=1) * (1.0000001 - true_probs))
+    loss = loss_adv + float(beta) * loss_robust
+
+
+        
+    return loss

robust_loss/rslad.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import torch.optim as optim
+import numpy as np
+
+# def attack_pgd(model,train_batch_data,train_batch_labels,attack_iters=10,step_size=2/255.0,epsilon=8.0/255.0):
+#     ce_loss = torch.nn.CrossEntropyLoss().cuda()
+#     train_ifgsm_data = train_batch_data.detach() + torch.zeros_like(train_batch_data).uniform_(-epsilon,epsilon)
+#     train_ifgsm_data = torch.clamp(train_ifgsm_data,0,1)
+#     for i in range(attack_iters):
+#         train_ifgsm_data.requires_grad_()
+#         logits = model(train_ifgsm_data)
+#         loss = ce_loss(logits,train_batch_labels.cuda())
+#         loss.backward()
+#         train_grad = train_ifgsm_data.grad.detach()
+#         train_ifgsm_data = train_ifgsm_data + step_size*torch.sign(train_grad)
+#         train_ifgsm_data = torch.clamp(train_ifgsm_data.detach(),0,1)
+#         train_ifgsm_pert = train_ifgsm_data - train_batch_data
+#         train_ifgsm_pert = torch.clamp(train_ifgsm_pert,-epsilon,epsilon)
+#         train_ifgsm_data = train_batch_data + train_ifgsm_pert
+#         train_ifgsm_data = train_ifgsm_data.detach()
+#     return train_ifgsm_data
+def kl_loss(a,b):
+    loss = -a*b + torch.log(b+1e-5)*b
+    return loss
+def rslad_inner_loss(model,
+                teacher_logits,
+                x_natural,
+                y,
+                optimizer,
+                step_size=0.0078,
+                epsilon=0.031,
+                perturb_steps=10,
+                beta=6.0):
+    # define KL-loss
+    criterion_kl = nn.KLDivLoss(size_average=False,reduce=False)
+    model.eval()
+    batch_size = len(x_natural)
+    # generate adversarial example
+    x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cuda().detach()
+
+    for _ in range(perturb_steps):
+        x_adv.requires_grad_()
+        with torch.enable_grad():
+            loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
+                                       F.softmax(teacher_logits, dim=1))
+            loss_kl = torch.sum(loss_kl)
+        grad = torch.autograd.grad(loss_kl, [x_adv])[0]
+        x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
+        x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)
+        x_adv = torch.clamp(x_adv, 0.0, 1.0)
+
+    model.train()
+
+    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
+    # zero gradient
+    # optimizer.zero_grad()
+    # logits = model(x_adv)
+    return x_adv
+def rslad_loss(teacher_model,model,x_natural,y,optimizer,step_size=0.0078,
+                epsilon=0.031,
+                perturb_steps=10,
+                beta=6.0):
+    teacher_logits = teacher_model(x_natural)
+    x_adv = rslad_inner_loss(model,teacher_logits,x_natural,y,optimizer,step_size,epsilon,perturb_steps)
+    adv_logits = model(x_adv)
+    model.train()
+    nat_logits = model(x_natural)
+    kl_Loss1 = kl_loss(F.log_softmax(adv_logits,dim=1),F.softmax(teacher_logits.detach(),dim=1))
+    kl_Loss2 = kl_loss(F.log_softmax(nat_logits,dim=1),F.softmax(teacher_logits.detach(),dim=1))
+    kl_Loss1 = torch.mean(kl_Loss1)
+    kl_Loss2 = torch.mean(kl_Loss2)
+    loss = 5/6.0*kl_Loss1 + 1/6.0*kl_Loss2
+    return loss

robust_loss/sat.py

@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from attacks import create_attack
+import numpy as np
+from torch.autograd import Variable
+from contextlib import contextmanager
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+else:
+    device = torch.device('cpu')
+
+class ctx_noparamgrad(object):
+    def __init__(self, module):
+        self.prev_grad_state = get_param_grad_state(module)
+        self.module = module
+        set_param_grad_off(module)
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):
+        set_param_grad_state(self.module, self.prev_grad_state)
+        return False
+
+
+class ctx_eval(object):
+    def __init__(self, module):
+        self.prev_training_state = get_module_training_state(module)
+        self.module = module
+        set_module_training_off(module)
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):
+        set_module_training_state(self.module, self.prev_training_state)
+        return False
+
+
+@contextmanager
+def ctx_noparamgrad_and_eval(module):
+    with ctx_noparamgrad(module) as a, ctx_eval(module) as b:
+        yield (a, b)
+
+
+def get_module_training_state(module):
+    return {mod: mod.training for mod in module.modules()}
+
+
+def set_module_training_state(module, training_state):
+    for mod in module.modules():
+        mod.training = training_state[mod]
+
+
+def set_module_training_off(module):
+    for mod in module.modules():
+        mod.training = False
+
+
+def get_param_grad_state(module):
+    return {param: param.requires_grad for param in module.parameters()}
+
+
+def set_param_grad_state(module, grad_state):
+    for param in module.parameters():
+        param.requires_grad = grad_state[param]
+
+
+def set_param_grad_off(module):
+    for param in module.parameters():
+        param.requires_grad = False
+class MadrysLoss(nn.Module):
+    def __init__(self, step_size=0.007, epsilon=0.031, perturb_steps=10, beta=6.0,
+                 distance='l_inf', cutmix=False, adjust_freeze=True, cutout=False,
+                 cutout_length=16):
+        super(MadrysLoss, self).__init__()
+        self.step_size = step_size
+        self.epsilon = epsilon
+        self.perturb_steps = perturb_steps
+        self.beta = beta
+        self.distance = distance
+        self.cross_entropy =  torch.nn.CrossEntropyLoss()
+        self.adjust_freeze = adjust_freeze
+        self.cutout = cutout
+        self.cutout_length = cutout_length
+
+    def forward(self, model, x_natural, labels): #optimizer
+        model.eval()
+        if self.adjust_freeze:
+            for param in model.parameters():
+                param.requires_grad = False
+
+        # generate adversarial example
+        x_adv = x_natural.detach() + self.step_size * torch.randn(x_natural.shape).to(device).detach()
+        if self.distance == 'l_inf':
+            adv_loss = 0
+            for _ in range(self.perturb_steps):
+                x_adv.requires_grad_()
+                loss_ce = self.cross_entropy(model(x_adv), labels)
+                grad = torch.autograd.grad(loss_ce, [x_adv])[0]
+                x_adv = x_adv.detach() + self.step_size * torch.sign(grad.detach())
+                x_adv = torch.min(torch.max(x_adv, x_natural - self.epsilon), x_natural + self.epsilon)
+                x_adv = torch.clamp(x_adv, 0.0, 1.0)
+        else:
+            x_adv = torch.clamp(x_adv, 0.0, 1.0)
+
+        x_adv = Variable(x_adv, requires_grad=False)
+
+        if self.adjust_freeze:
+            for param in model.parameters():
+                param.requires_grad = True
+
+        if self.cutout:
+            batch_size = x_adv.shape[0]
+            c, h, w = x_adv.shape[1], x_adv.shape[2], x_adv.shape[3]
+            mask = torch.ones(batch_size, c, h, w).float()
+            for j in range(batch_size):
+                y = np.random.randint(h)
+                x = np.random.randint(w)
+
+                y1 = np.clip(y - self.cutout_length // 2, 0, h)
+                y2 = np.clip(y + self.cutout_length // 2, 0, h)
+                x1 = np.clip(x - self.cutout_length // 2, 0, w)
+                x2 = np.clip(x + self.cutout_length // 2, 0, w)
+
+                mask[j, :, y1: y2, x1: x2] = 0.0
+            x_adv = x_adv * mask.to(device)
+
+        model.train()
+        # optimizer.zero_grad()
+
+
+        logits = model(x_adv)
+        loss = self.cross_entropy(logits, labels)
+
+        return loss
+
+
+def sat_loss(model, x, y,optimizer,step_size,epsilon,num_steps,attack_type,beta,criterion= torch.nn.CrossEntropyLoss()):
+    """
+    Adversarial training (Madry et al, 2017).
+    """
+    attack = create_attack(model, criterion, 'linf-pgd', epsilon, num_steps, step_size)  
+    with ctx_noparamgrad_and_eval(model):
+        x_adv, _ = attack.perturb(x, y)
+    print(x_adv.shape)
+    y_adv = y
+    out = model(x_adv)
+    loss = criterion(out, y_adv)
+    
+    return loss  
+
+
+
+
+
+

robust_loss/trades.py

@@ -0,0 +1,92 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import torch.optim as optim
+from utils.criterion import CrossEntropyWithLabelSmooth
+
+
+
+def squared_l2_norm(x):
+    flattened = x.view(x.unsqueeze(0).shape[0], -1)
+    return (flattened ** 2).sum(1)
+
+
+def l2_norm(x):
+    return squared_l2_norm(x).sqrt()
+
+
+def trades_loss(model, x_natural, y,optimizer = None, step_size=0.003, epsilon=0.031, perturb_steps=10, beta=1.0, 
+                attack='l_inf',natural_criterion= nn.CrossEntropyLoss() ):
+    """
+    TRADES training (Zhang et al, 2019).
+    """
+  
+    # define KL-loss
+    criterion_kl = nn.KLDivLoss(size_average=False)
+    model.eval()
+    batch_size = len(x_natural)
+    # generate adversarial example
+    x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).cuda().detach()
+    p_natural = F.softmax(model(x_natural), dim=1)
+    
+    if attack == 'l_inf':
+        for _ in range(perturb_steps):
+            x_adv.requires_grad_()
+            with torch.enable_grad():
+                loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1), p_natural)
+            grad = torch.autograd.grad(loss_kl, [x_adv])[0]
+            x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
+            x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)
+            x_adv = torch.clamp(x_adv, 0.0, 1.0)
+    
+    elif attack == 'l2':
+        delta = 0.001 * torch.randn(x_natural.shape).cuda().detach()
+        delta = Variable(delta.data, requires_grad=True)
+
+        # Setup optimizers
+        optimizer_delta = optim.SGD([delta], lr=epsilon / perturb_steps * 2)
+
+        for _ in range(perturb_steps):
+            adv = x_natural + delta
+
+            # optimize
+            optimizer_delta.zero_grad()
+            with torch.enable_grad():
+                loss = (-1) * criterion_kl(F.log_softmax(model(adv), dim=1), p_natural)
+            loss.backward(retain_graph=True)
+            # renorming gradient
+            grad_norms = delta.grad.view(batch_size, -1).norm(p=2, dim=1)
+            delta.grad.div_(grad_norms.view(-1, 1, 1, 1))
+            # avoid nan or inf if gradient is 0
+            if (grad_norms == 0).any():
+                delta.grad[grad_norms == 0] = torch.randn_like(delta.grad[grad_norms == 0])
+            optimizer_delta.step()
+
+            # projection
+            delta.data.add_(x_natural)
+            delta.data.clamp_(0, 1).sub_(x_natural)
+            delta.data.renorm_(p=2, dim=0, maxnorm=epsilon)
+        x_adv = Variable(x_natural + delta, requires_grad=False)
+    else:
+        raise ValueError(f'Attack={attack} not supported for TRADES training!')
+    model.train()
+
+    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
+    
+    optimizer.zero_grad()
+    # calculate robust loss
+    logits_natural = model(x_natural)
+    # print("loguts natural:{}".format(logits_natural))
+    logits_adv = model(x_adv)
+    # print("loguts adv:{}".format(logits_adv))
+    loss_natural = natural_criterion(logits_natural, y)
+    loss_robust = (1.0 / batch_size) * criterion_kl(F.log_softmax(logits_adv, dim=1),
+                                                    F.softmax(logits_natural, dim=1))
+    loss = loss_natural + beta * loss_robust
+    
+    
+        
+    return loss