This repository has been archived by the owner on Sep 12, 2023. It is now read-only.
/
make_simulation.py
83 lines (63 loc) · 3.25 KB
/
make_simulation.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
from common import ContrastiveNetwork,train_contrastive_network, get_embs, fit_surrogate, predict_surrogate, sim_one_run, embs_and_accs_function
import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from einops import rearrange, reduce, repeat
from tqdm import tqdm
import torch.nn as nn
import argparse
import SearchSpaces
sns.set()
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
parser = argparse.ArgumentParser(description='Generate prediction plot based on subset of samples')
parser.add_argument('path_EPDJMs', type=str)
parser.add_argument('task', choices=['CIFAR10', 'CIFAR100', 'ImageNet16'])
parser.add_argument('search_space', choices=['201', 'nats_ss'])
parser.add_argument('--num_simulations', type=int, default = 500)
parser.add_argument('--evaluation_budget', type=int, default = 200)
parser.add_argument('--lengthscale', type = float, default =1)
parser.add_argument('--noise', type = float, default =0.1)
parser.add_argument('--ei_offset', type = float, default =0.15)
args = parser.parse_args()
class CachedProjectedJacobians(Dataset):
def __init__(self, path, num_augs):
self.data = np.load(path, mmap_mode='r')
self.data_augs = self.data[0].shape[0]
self.proj_size = self.data[0].shape[2]
self.num_augs = num_augs
assert(num_augs <= self.data_augs)
def __len__(self):
return len(self.data)
def __getitem__(self, index):
return self.data[index, np.random.choice(self.data_augs, self.num_augs, replace=False)]
def sns_lineplot(x, y, *args, **kwargs): #sns doesn't broadcast 1D x to match 2D y, so we do that here.
l = len(x)
assert(y.shape[-1] == l)
y.reshape(-1, l)
n = len(y)
x = repeat(np.array(x), "x -> (b x)", b = n)
sns.lineplot(x, y.flatten(), *args, **kwargs)
data_set = CachedProjectedJacobians(path = args.path_EPDJMs, num_augs =2)
net = ContrastiveNetwork(data_set[0].shape[-1], emb_size = 256, projection_head_out_size=256, dim_hidden = 256)
net.cuda()
train_contrastive_network(net, data_set, batch_size=1024, epochs=30)
data_set = CachedProjectedJacobians(path = args.path_EPDJMs, num_augs = 4)
embs = get_embs(net,data_set)
validation = SearchSpaces.get_metrics(args.search_space, args.task)
plt.figure(figsize=(6,5))
data_func = embs_and_accs_function(embs, validation)
runs = [sim_one_run(data_func, args.evaluation_budget, num_randoms = 5, num_augs = 4, lengthscale = args.lengthscale, noise = args.noise, ei_offset = args.ei_offset) for _ in tqdm(range(args.num_simulations))]
runs = np.maximum.accumulate(runs, axis = 1)
random_runs = [sim_one_run(data_func, args.evaluation_budget, num_randoms = args.evaluation_budget, num_augs = 1) for _ in tqdm(range(args.num_simulations))]
random_runs = np.maximum.accumulate(random_runs, axis = 1)
sns_lineplot(np.arange(args.evaluation_budget), runs, label = 'Contrasive Embeddings + GP')
sns_lineplot(np.arange(args.evaluation_budget), random_runs, label = 'Random Search')
plt.ylim(bottom = random_runs[:, 5].mean())
plt.xlabel("Number of Evaluated Architectures")
plt.ylabel("Accuracy of Best Architecture")
plt.savefig("simulations"+args.task + ".pdf")
np.save("simulations" + args.task + ".npy", (runs, random_runs))