Upload files

Author: manyana72

.gitignore

@@ -0,0 +1,15 @@
+__pycache__
+tensorboard
+commands
+
+*.log
+*.pt
+*.tar
+*.pkl
+*.bat
+*.pth
+*.png
+*.jpg
+*.sh
+*.pdf
+*.info

LS.py

@@ -0,0 +1,19 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+class LabelSmoothingCrossEntropy(nn.Module):
+    def __init__(self, eps=0.1, reduction='mean'):
+        super(LabelSmoothingCrossEntropy, self).__init__()
+        self.eps = eps
+        self.reduction = reduction
+
+    def forward(self, output, target):
+        c = output.size()[-1]
+        log_preds = F.log_softmax(output, dim=-1)
+        if self.reduction=='sum':
+            loss = -log_preds.sum()
+        else:
+            loss = -log_preds.sum(dim=-1)
+            if self.reduction=='mean':
+                loss = loss.mean()
+        return loss*self.eps/c + (1-self.eps) * F.nll_loss(log_preds, target, reduction=self.reduction)

SAM.py

@@ -0,0 +1,63 @@
+import torch
+
+
+class SAM(torch.optim.Optimizer):
+    def __init__(self, params, base_optimizer, rho=0.05, adaptive=False, **kwargs):
+        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
+        super(SAM, self).__init__(params, defaults)
+
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+        self.defaults.update(self.base_optimizer.defaults)
+
+    @torch.no_grad()
+    def first_step(self, zero_grad=False):
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group["rho"] / (grad_norm + 1e-12)
+
+            for p in group["params"]:
+                if p.grad is None: continue
+                self.state[p]["old_p"] = p.data.clone()
+                e_w = (torch.pow(p, 2) if group["adaptive"] else 1.0) * p.grad * scale.to(p)
+                p.add_(e_w)  # climb to the local maximum "w + e(w)"
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad=False):
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None: continue
+                p.data = self.state[p]["old_p"]  # get back to "w" from "w + e(w)"
+
+        self.base_optimizer.step()  # do the actual "sharpness-aware" update
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        assert closure is not None, "Sharpness Aware Minimization requires closure, but it was not provided"
+        closure = torch.enable_grad()(closure)  # the closure should do a full forward-backward pass
+
+        self.first_step(zero_grad=True)
+        closure()
+        self.second_step()
+
+    def _grad_norm(self):
+        shared_device = self.param_groups[0]["params"][0].device  # put everything on the same device, in case of model parallelism
+        norm = torch.norm(
+                    torch.stack([
+                        ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
+                        for group in self.param_groups for p in group["params"]
+                        if p.grad is not None
+                    ]),
+                    p=2
+               )
+        return norm
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        self.base_optimizer.param_groups = self.param_groups

arg_parser.py

@@ -0,0 +1,104 @@
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='PyTorch Lottery Tickets Experiments')
+
+    ##################################### Dataset #################################################
+    parser.add_argument('--data', type=str, default='../data',
+                        help='location of the data corpus')
+    parser.add_argument('--dataset', type=str,
+                        default='cifar10', help='dataset')
+    parser.add_argument('--input_size', type=int,
+                        default=32, help='size of input images')
+    parser.add_argument('--data_dir', type=str,
+                        default='./tiny-imagenet-200', help='dir to tiny-imagenet')
+    parser.add_argument('--num_workers', type=int, default=4)
+    parser.add_argument('--num_classes', type=int, default=10)
+    ##################################### Architecture ############################################
+    parser.add_argument('--arch', type=str,
+                        default='resnet18', help='model architecture')
+    parser.add_argument('--imagenet_arch', action="store_true",
+                        help="architecture for imagenet size samples")
+    parser.add_argument('--train_y_file', type=str,
+                        default='./labels/train_ys.pth', help='labels for training files')
+    parser.add_argument('--val_y_file', type=str,
+                        default='./labels/val_ys.pth', help='labels for validation files')
+    ##################################### General setting ############################################
+    parser.add_argument('--seed', default=2, type=int, help='random seed')
+    parser.add_argument('--train_seed', default=1, type=int,
+                        help='seed for training (default value same as args.seed)')
+    parser.add_argument('--gpu', type=int, default=0, help='gpu device id')
+    parser.add_argument('--workers', type=int, default=4,
+                        help='number of workers in dataloader')
+    parser.add_argument('--resume', action="store_true",
+                        help="resume from checkpoint")
+    parser.add_argument('--checkpoint', type=str,
+                        default=None, help='checkpoint file')
+    parser.add_argument(
+        '--save_dir', help='The directory used to save the trained models', default=None, type=str)
+    parser.add_argument('--mask', type=str, default=None, help='sparse model')
+
+    ##################################### Training setting #################################################
+    parser.add_argument('--batch_size', type=int,
+                        default=256, help='batch size')
+    parser.add_argument('--lr', default=0.1, type=float,
+                        help='initial learning rate')
+    parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
+    parser.add_argument('--weight_decay', default=5e-4,
+                        type=float, help='weight decay')
+    parser.add_argument('--epochs', default=182, type=int,
+                        help='number of total epochs to run')
+    parser.add_argument('--warmup', default=0, type=int, help='warm up epochs')
+    parser.add_argument('--print_freq', default=50,
+                        type=int, help='print frequency')
+    parser.add_argument('--decreasing_lr', default='91,136',
+                        help='decreasing strategy')
+    parser.add_argument('--no-aug', action='store_true', default=False,
+                        help='No augmentation in training dataset (transformation).')
+    parser.add_argument('--no-l1-epochs', default=0, type=int, help='non l1 epochs')
+    ##################################### Pruning setting #################################################
+    parser.add_argument('--prune', type=str, default="omp",
+                        help="method to prune")
+    parser.add_argument('--pruning_times', default=1,
+                        type=int, help='overall times of pruning (only works for IMP)')
+    parser.add_argument('--rate', default=0.95, type=float,
+                        help='pruning rate')  # pruning rate is always 20%
+    parser.add_argument('--prune_type', default='rewind_lt', type=str,
+                        help='IMP type (lt, pt or rewind_lt)')
+    parser.add_argument('--random_prune', action='store_true',
+                        help='whether using random prune')
+    parser.add_argument('--rewind_epoch', default=0,
+                        type=int, help='rewind checkpoint')
+    parser.add_argument('--rewind_pth', default=None,
+                        type=str, help='rewind checkpoint to load')
+    parser.add_argument('--hf_vit', default='NO', type=str, 
+                        choices=['YES', 'NO'], help='lora method')
+    parser.add_argument('--debug', default=False, type=bool, 
+                        choices=[True, False], help='lora method')
+
+    ##################################### Unlearn setting #################################################
+    parser.add_argument('--unlearn', type=str,
+                        default='retrain', help='method to unlearn')
+    parser.add_argument('--unlearn_lr', default=0.01, type=float,
+                        help='initial learning rate')
+    parser.add_argument('--unlearn_epochs', default=10, type=int,
+                        help='number of total epochs for unlearn to run')
+    parser.add_argument('--num_indexes_to_replace', type=int, default=None,
+                        help='Number of data to forget')
+    parser.add_argument('--class_to_replace', type=int, default=0,
+                        help='Specific class to forget')
+    parser.add_argument('--indexes_to_replace', type=list, default=None,
+                        help='Specific index data to forget')
+    parser.add_argument('--alpha', default=0.2, type=float,
+                        help='unlearn noise')
+    parser.add_argument('--lora', default='NO', type=str, 
+                        choices=['YES', 'NO'], help='lora method')
+
+    ##################################### Attack setting #################################################
+    parser.add_argument('--attack', type=str,
+                        default='backdoor', help='method to unlearn')
+    parser.add_argument('--trigger_size', type=int, default=4,
+                        help='The size of trigger of backdoor attack')
+    return parser.parse_args()