-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathimagenet.py
135 lines (124 loc) · 6.38 KB
/
imagenet.py
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import torchvision
from torch.utils.data import DataLoader, Subset
from datasets.load import load_dataset
import os
import sys
# sys.path.append(".")
# from cfg import *
from tqdm import tqdm
import torch
def prepare_data(dataset, batch_size = 512, shuffle=True, train_subset_indices=None, val_subset_indices=None,data_path='/localscratch/dataset'):
path = os.path.join(data_path, "huggingface")
if dataset == "imagenet":
train_set = load_dataset("imagenet-1k", use_auth_token=True, split="train", cache_dir=path)
validation_set = load_dataset("imagenet-1k", use_auth_token=True, split="validation", cache_dir=path)
def train_transform(examples):
transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda x: x.convert('RGB')),
torchvision.transforms.RandomResizedCrop((224, 224)),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
])
examples["image"] = [transform(x) for x in examples["image"]]
return examples
def validation_transform(examples):
transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda x: x.convert('RGB')),
torchvision.transforms.Resize((256, 256)),
torchvision.transforms.CenterCrop((224, 224)),
torchvision.transforms.ToTensor(),
])
examples["image"] = [transform(x) for x in examples["image"]]
return examples
elif dataset == "tiny_imagenet":
train_set = load_dataset("Maysee/tiny-imagenet", use_auth_token=True, split="train", cache_dir=path)
validation_set = load_dataset("Maysee/tiny-imagenet", use_auth_token=True, split="valid", cache_dir=path)
def train_transform(examples):
transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda x: x.convert('RGB')),
torchvision.transforms.RandomCrop(64, padding=4),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
examples["image"] = [transform(x) for x in examples["image"]]
return examples
def validation_transform(examples):
transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda x: x.convert('RGB')),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
examples["image"] = [transform(x) for x in examples["image"]]
return examples
elif dataset == "flowers102":
train_set = load_dataset("nelorth/oxford-flowers", use_auth_token=True, split="train", cache_dir=path)
validation_set = load_dataset("nelorth/oxford-flowers", use_auth_token=True, split="test", cache_dir=path)
def train_transform(examples):
transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda x: x.convert('RGB')),
torchvision.transforms.Resize((256, 256)),
torchvision.transforms.RandomCrop((224, 224)),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
examples["image"] = [transform(x) for x in examples["image"]]
return examples
def validation_transform(examples):
transform = torchvision.transforms.Compose([
torchvision.transforms.Lambda(lambda x: x.convert('RGB')),
torchvision.transforms.Resize((256, 256)),
torchvision.transforms.CenterCrop((224, 224)),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
examples["image"] = [transform(x) for x in examples["image"]]
return examples
else:
raise NotImplementedError
train_set.set_transform(transform=train_transform)
validation_set.set_transform(transform=validation_transform)
if train_subset_indices is not None:
forget_indices = torch.ones_like(train_subset_indices)-train_subset_indices
train_subset_indices = torch.nonzero(train_subset_indices)
forget_indices = torch.nonzero(forget_indices)
retain_set = Subset(train_set, train_subset_indices)
forget_set = Subset(train_set,forget_indices)
if val_subset_indices is not None:
val_subset_indices = torch.nonzero(val_subset_indices)
validation_set = Subset(validation_set, val_subset_indices)
if train_subset_indices is not None:
loaders = {
'train': DataLoader(retain_set, batch_size = batch_size, num_workers = 12, shuffle = shuffle),
'val': DataLoader(validation_set, batch_size = batch_size, num_workers = 12, shuffle = shuffle),
'fog': DataLoader(forget_set, batch_size = batch_size, num_workers = 12, shuffle = shuffle)
}
else:
loaders = {
'train': DataLoader(train_set, batch_size = batch_size, num_workers = 12, shuffle = shuffle),
'val': DataLoader(validation_set, batch_size = batch_size, num_workers = 12, shuffle = shuffle),
}
return loaders
def get_x_y_from_data_dict(data, device):
x, y = data.values()
if isinstance(x, list):
x, y = x[0].to(device), y[0].to(device)
else:
x, y = x.to(device), y.to(device)
return x, y
if __name__ == "__main__":
ys = {}
ys['train'] = []
ys['val'] = []
loaders = prepare_data(dataset='imagenet', batch_size=1, shuffle=False)
for data in tqdm(loaders['val'], ncols=100):
x, y = get_x_y_from_data_dict(data, 'cpu')
ys['val'].append(y.item())
for data in tqdm(loaders['train'], ncols=100):
x, y = get_x_y_from_data_dict(data, 'cpu')
ys['train'].append(y.item())
ys['train'] = torch.Tensor(ys['train']).long()
ys['val'] = torch.Tensor(ys['val']).long()
torch.save(ys['train'], 'train_ys.pth')
torch.save(ys['val'],'val_ys.pth')