MachineUnlearning-DocClassification/dataset_unlearn_aug.py at main · leitro/MachineUnlearning-DocClassification · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import os
import random
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
import torch
from torchvision.transforms import v2

IMG = '/data/users/lkang/RVL-CDIP/images/'
LABEL_F = '/data/users/lkang/RVL-CDIP/RVL_CDIP_full.npy'

'''Random sampling'''
LABEL_S = '/data/users/lkang/RVL-CDIP/RVL_CDIP_subset_0.1.npy'
#LABEL_S = '/data/users/lkang/RVL-CDIP/RVL_CDIP_subset_0.05.npy'
#LABEL_S = '/data/users/lkang/RVL-CDIP/RVL_CDIP_subset_0.01.npy'

NUM_CLASS = 16

UNCLASS = [0] # unlearn classes


class RVL(Dataset):
    '''retain: True return retain set, False return unlearn set'''
    def __init__(self, img_dir, label_dir, unlearn_classes, split, retain_o_forget, random_label=False): # split: train, valid, test
        self.random_label = random_label
        data_all = np.load(label_dir, allow_pickle=True).item()
        self.data = self._retain_unlearn(data_all[split], unlearn_classes, retain_o_forget)
        self.img_proc = torch.nn.Sequential(
            v2.ToImage(),
            v2.ToDtype(torch.float32, scale=True),
            v2.Resize((224, 224), antialias=True),
            v2.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        )
        self.img_proc_aug = torch.nn.Sequential(
            v2.ToImage(),
            v2.ToDtype(torch.float32, scale=True),
            v2.Resize((224, 224), antialias=True),
            v2.RandomResizedCrop((224, 224), scale=(0.7, 1), ratio=(0.85, 1.15), antialias=True),
            v2.RandomHorizontalFlip(p=0.4),
            v2.GaussianBlur(kernel_size=5, sigma=(0.1, 5)),
            v2.RandomAdjustSharpness(sharpness_factor=2, p=0.4),
            v2.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
        )

    def _retain_unlearn(self, data, unlearn_classes, retain_o_forget): # return: retain_set or forget_set
        if retain_o_forget: # True: retain_set
            return [item for item in data if int(item[1]) not in unlearn_classes]
        else: # False: forget_set
            return [item for item in data if int(item[1]) in unlearn_classes]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        record = self.data[idx]
        img_url = record[0]
        if self.random_label:
            ids = list(range(NUM_CLASS))
            for cla in UNCLASS:
                ids.remove(cla)
            gt = random.choice(ids)
        else:
            clas = record[1]
            gt = int(clas)

        img = Image.open(f'{IMG}{img_url}').convert('RGB')
        img_feat = self.img_proc(img)


        sample_info = {'img_url': img_url,
                       'img': img_feat,
                       'label': gt,
                       }

        return sample_info


def loadData_full():
    data_dir = dict()
    for split in ['train', 'valid', 'test']:
        data_dir[split] = dict()
        data_dir[split]['retain'] = RVL(IMG, LABEL_F, UNCLASS, split, True)
        data_dir[split]['forget'] = RVL(IMG, LABEL_F, UNCLASS, split, False)
    return data_dir


def loadData_subset():
    data_dir = dict()
    for split in ['train', 'valid', 'test']:
        data_dir[split] = dict()
        data_dir[split]['retain'] = RVL(IMG, LABEL_S, UNCLASS, split, True)
        data_dir[split]['forget'] = RVL(IMG, LABEL_S, UNCLASS, split, False)
        data_dir[split]['forget_randomlabel'] = RVL(IMG, LABEL_S, UNCLASS, split, False, True)
    return data_dir


def loadData_subset_contour(LABEL_SC):
    split = 'train'
    data_dir = dict()
    data_dir[split] = dict()
    data_dir[split]['retain'] = RVL(IMG, LABEL_SC, UNCLASS, split, True)
    data_dir[split]['forget'] = RVL(IMG, LABEL_SC, UNCLASS, split, False)
    data_dir[split]['forget_randomlabel'] = RVL(IMG, LABEL_SC, UNCLASS, split, False, True)
    return data_dir


if __name__ == '__main__':
    res = loadData()
    next(iter(res['train']['forget']))