Create scrub.py

scrub.py

@@ -0,0 +1,358 @@
+import argparse
+import json
+import os
+import time
+import copy
+import random
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.optim.lr_scheduler import LinearLR
+from sklearn.linear_model import LogisticRegression
+from IPython import embed
+import torchvision.models as models
+import matplotlib.pyplot as plt
+import datetime
+import torch.nn as nn
+from transformers import ViTModel, ViTFeatureExtractor
+import lime
+from lime import lime_image
+from skimage.segmentation import mark_boundaries
+from tqdm import tqdm
+
+import os, torch, shutil, numpy as np
+from glob import glob; from PIL import Image
+from torch.utils.data import random_split, Dataset, DataLoader
+from torchvision import transforms as T
+torch.manual_seed(2024)
+
+import sys
+import time
+from torch import nn
+from itertools import cycle
+import timm
+
+
+class CustomDataset(Dataset):
+    
+    def __init__(self, root, transformations = None):
+        
+        self.transformations = transformations
+        self.im_paths = [im_path for im_path in sorted(glob(f"{root}/*/*"))]
+        self.im_paths = [i for i in self.im_paths if not 'Will Smith' in i]
+
+        self.cls_names, self.cls_counts, count, data_count = {}, {}, 0, 0
+        for idx, im_path in enumerate(self.im_paths):
+            class_name = self.get_class(im_path)
+            if class_name not in self.cls_names: self.cls_names[class_name] = count; self.cls_counts[class_name] = 1; count += 1
+            else: self.cls_counts[class_name] += 1
+        
+    def get_class(self, path): return os.path.dirname(path).split("/")[-1]
+    
+    def __len__(self): return len(self.im_paths)
+
+    def __getitem__(self, idx):
+        
+        im_path = self.im_paths[idx]
+        im = Image.open(im_path).convert("RGB")
+        gt = self.cls_names[self.get_class(im_path)]
+        
+        if self.transformations is not None: im = self.transformations(im)
+        
+        return im, gt
+
+
+class SingleCelebCustomDataset(Dataset):
+    
+    def __init__(self, root, transformations = None):
+        
+        self.transformations = transformations
+        self.im_paths = [im_path for im_path in sorted(glob(f"{root}/*"))]
+        self.cls_names, self.cls_counts, count, data_count = {}, {}, 0, 0
+        for idx, im_path in enumerate(self.im_paths):
+            class_name = self.get_class(im_path)
+            if class_name not in self.cls_names: self.cls_names[class_name] = count; self.cls_counts[class_name] = 1; count += 1
+            else: self.cls_counts[class_name] += 1
+        
+    def get_class(self, path): return 16
+    
+    def __len__(self): return len(self.im_paths)
+
+    def __getitem__(self, idx):
+        
+        im_path = self.im_paths[idx]
+        im = Image.open(im_path).convert("RGB")
+        gt = self.cls_names[self.get_class(im_path)]
+        
+        if self.transformations is not None: im = self.transformations(im)
+        
+        return im, gt
+
+    
+def get_dls(root, transformations, bs, split = [0.9, 0.05, 0.05], ns = 4, single=False):
+    if single:
+        ds = SingleCelebCustomDataset(root = root, transformations = transformations)
+    else: 
+        ds = CustomDataset(root = root, transformations = transformations)
+
+    total_len = len(ds)
+    tr_len = int(total_len * split[0])
+    vl_len = int(total_len * split[1])
+    ts_len = total_len - (tr_len + vl_len)
+    
+    tr_ds, vl_ds, ts_ds = random_split(dataset = ds, lengths = [tr_len, vl_len, ts_len])
+    
+    tr_dl, val_dl, ts_dl = DataLoader(tr_ds, batch_size = bs, shuffle = True, num_workers = ns), DataLoader(vl_ds, batch_size = bs, shuffle = False, num_workers = ns), DataLoader(ts_ds, batch_size = 1, shuffle = False, num_workers = ns)
+    
+    return tr_dl, val_dl, ts_dl, ds.cls_names
+
+
+def param_dist(model, swa_model, p):
+    #This is from https://github.com/ojus1/SmoothedGradientDescentAscent/blob/main/SGDA.py
+    dist = 0.
+    for p1, p2 in zip(model.parameters(), swa_model.parameters()):
+        dist += torch.norm(p1 - p2, p='fro')
+    return p * dist
+
+def adjust_learning_rate_new(epoch, optimizer, LUT):
+    """
+    new learning rate schedule according to RotNet
+    """
+    lr = next((lr for (max_epoch, lr) in LUT if max_epoch > epoch), LUT[-1][1])
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def sgda_adjust_learning_rate(epoch, opt, optimizer):
+    """Sets the learning rate to the initial LR decayed by decay rate every steep step"""
+    steps = np.sum(epoch > np.asarray(opt.lr_decay_epochs))
+    new_lr = opt.sgda_learning_rate
+    if steps > 0:
+        new_lr = opt.sgda_learning_rate * (opt.lr_decay_rate ** steps)
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = new_lr
+    return new_lr
+    
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+
+
+def train_distill(epoch, train_loader, module_list, swa_model, criterion_list, optimizer, opt, split, quiet=False):
+    """One epoch distillation"""
+    # set modules as train()
+    for module in module_list:
+        module.train()
+    # set teacher as eval()
+    module_list[-1].eval()
+
+
+    criterion_cls = criterion_list[0]
+    criterion_div = criterion_list[1]
+    criterion_kd = criterion_list[2]
+
+    model_s = module_list[0]
+    model_t = module_list[-1]
+
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses = AverageMeter()
+    kd_losses = AverageMeter()
+    top1 = AverageMeter()
+
+
+    end = time.time()
+    for idx, data in enumerate(train_loader):
+        if opt.distill in ['crd']:
+            input, target, index, contrast_idx = data
+        else:
+            input, target = data
+        data_time.update(time.time() - end)
+
+        input = input.float()
+        if torch.cuda.is_available():
+            input = input.cuda()
+            target = target.cuda()
+            if opt.distill in ['crd']:
+                contrast_idx = contrast_idx.cuda()
+                index = index.cuda()
+
+        # ===================forward=====================
+        #feat_s, logit_s = model_s(input, is_feat=True, preact=False)
+        logit_s = model_s(input)
+        with torch.no_grad():
+            #feat_t, logit_t = model_t(input, is_feat=True, preact=preact)
+            #feat_t = [f.detach() for f in feat_t]
+            logit_t = model_t(input)
+
+
+        # cls + kl div
+        loss_cls = criterion_cls(logit_s, target)
+        loss_div = criterion_div(logit_s, logit_t)
+
+        if split == "minimize":
+            loss = opt.gamma * loss_cls + opt.alpha * loss_div
+        elif split == "maximize":
+            loss = -loss_div
+
+        loss = loss + param_dist(model_s, swa_model, opt.smoothing)
+
+        if split == "minimize" and not quiet:
+            acc1, _ = accuracy(logit_s, target, topk=(1,1))
+            losses.update(val=loss.item(), n=input.size(0))
+            top1.update(val=acc1[0], n=input.size(0))
+        elif split == "maximize" and not quiet:
+            kd_losses.update(val=loss.item(), n=input.size(0))
+        elif split == "linear" and not quiet:
+            acc1, _ = accuracy(logit_s, target, topk=(1, 1))
+            losses.update(val=loss.item(), n=input.size(0))
+            top1.update(val=acc1[0], n=input.size(0))
+            kd_losses.update(val=loss.item(), n=input.size(0))
+
+        # ===================backward=====================
+        optimizer.zero_grad()
+        loss.backward()
+        #nn.utils.clip_grad_value_(model_s.parameters(), clip)
+        optimizer.step()
+
+        # ===================meters=====================
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if not quiet:
+            if split == "mainimize":
+                if idx % opt.print_freq == 0:
+                    print('Epoch: [{0}][{1}/{2}]\t'
+                          'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                          'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
+                          'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
+                          'Acc@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
+                        epoch, idx, len(train_loader), batch_time=batch_time,
+                        data_time=data_time, loss=losses, top1=top1))
+                    sys.stdout.flush()
+
+    
+    if split == "minimize":
+        #if not quiet:
+            #print(' * Acc@1 {top1.avg:.3f} '
+             #     .format(top1=top1))
+
+        return top1.avg, losses.avg
+    else:
+        return kd_losses.avg
+
+
+class DistillKL(nn.Module):
+    """Distilling the Knowledge in a Neural Network"""
+    def __init__(self, T):
+        super(DistillKL, self).__init__()
+        self.T = T
+
+    def forward(self, y_s, y_t):
+        p_s = F.log_softmax(y_s/self.T, dim=1)
+        p_t = F.softmax(y_t/self.T, dim=1)
+        loss = F.kl_div(p_s, p_t, size_average=False) * (self.T**2) / y_s.shape[0]
+        return loss
+
+class Args:
+    def __init__(self, **entries):
+        self.__dict__.update(entries)
+
+def unlearn(): 
+    will_tr_dl, will_val_dl, will_ts_dl, classes = get_dls(root = "forget_set/", transformations = tfs, bs = 32, single=True)
+    model = timm.create_model("rexnet_150", pretrained = True, num_classes = 17)  
+    model.load_state_dict(torch.load('faces_best_model.pth'))
+    args = Args()
+    args.optim = 'sgd'
+    args.gamma = 0.99
+    args.alpha = 0.001
+    args.smoothing = 0.0
+    args.msteps = 4
+    args.clip = 0.2
+    args.sstart = 10
+    args.kd_T = 4
+    args.distill = 'kd'
+    
+    args.sgda_batch_size = 64
+    args.del_batch_size = 64
+    args.sgda_epochs = 6
+    args.sgda_learning_rate = 0.005
+    args.lr_decay_epochs = [3,5,9]
+    args.lr_decay_rate = 0.0005
+    args.sgda_weight_decay = 5e-4
+    args.sgda_momentum = 0.9
+    model_t = copy.deepcopy(model)
+    model_s = copy.deepcopy(model)
+    swa_model = torch.optim.swa_utils.AveragedModel(
+    model_s, avg_fn=avg_fn)
+    module_list = nn.ModuleList([])
+    module_list.append(model_s)
+    trainable_list = nn.ModuleList([])
+    trainable_list.append(model_s)
+    
+    criterion_cls = nn.CrossEntropyLoss()
+    criterion_div = DistillKL(args.kd_T)
+    criterion_kd = DistillKL(args.kd_T)
+    
+    
+    criterion_list = nn.ModuleList([])
+    criterion_list.append(criterion_cls)    # classification loss
+    criterion_list.append(criterion_div)    # KL divergence loss, original knowledge distillation
+    criterion_list.append(criterion_kd)     # other knowledge distillation loss
+    if args.optim == "sgd":
+        optimizer = optim.SGD(trainable_list.parameters(),
+                              lr=args.sgda_learning_rate,
+                              momentum=args.sgda_momentum,
+                              weight_decay=args.sgda_weight_decay)
+    
+    module_list.append(model_t)
+    
+    if torch.cuda.is_available():
+        module_list.cuda()
+        criterion_list.cuda()
+        import torch.backends.cudnn as cudnn
+        cudnn.benchmark = True
+        swa_model.cuda()
+    for epoch in tqdm(range(1, args.sgda_epochs + 1)):
+        maximize_loss = 0
+        if epoch <= args.msteps:
+            maximize_loss = train_distill(epoch, will_tr_dl, module_list, swa_model, criterion_list, optimizer, args, "maximize")
+        train_acc, train_loss = train_distill(epoch, celebs_tr_dl, module_list, swa_model, criterion_list, optimizer, args, "minimize")
+        if epoch >= args.sstart:
+            swa_model.update_parameters(model_s)
+