-
Notifications
You must be signed in to change notification settings - Fork 3
/
toy.py
executable file
·157 lines (108 loc) · 6.57 KB
/
toy.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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import hydra
from time import time
from pathlib import Path
import numpy as np
import torch
from torch.optim import Adam
from core.bounds import BOUNDS
from core.losses import sigmoid_loss, moment_loss, rand_loss
from core.monitors import MonitorMV
from core.utils import deterministic
from data.datasets import Dataset
from models.majority_vote import MultipleMajorityVote, MajorityVote
from models.random_forest import two_forests
from models.stumps import uniform_decision_stumps, custom_decision_stumps
from optimization import train_batch
from graphics.plot_predictions import plot_2D
import matplotlib.pyplot as plt
@hydra.main(config_path='config/toy.yaml')
def main(cfg):
SAVE_DIR = f"{hydra.utils.get_original_cwd()}/results/{cfg.dataset.distr}/noise={cfg.dataset.noise}/{cfg.dataset.N_train}/{cfg.training.risk}/{cfg.bound.type}/optimize-bound={cfg.training.opt_bound}/{cfg.model.pred}/M={cfg.model.M}/prior={cfg.model.prior}/lr={cfg.training.lr}/seeds={cfg.training.seed}-{int(cfg.training.seed)+cfg.num_trials}/"
SAVE_DIR = Path(SAVE_DIR)
if cfg.training.risk == "MC":
SAVE_DIR = SAVE_DIR / f"MC={cfg.training.MC_draws}"
if (SAVE_DIR / "err-b.npy").is_file():
print("results already saved in", SAVE_DIR)
exit(0)
SAVE_DIR.mkdir(parents=True, exist_ok=True)
print("results will be saved in:", SAVE_DIR.resolve())
# define params for each method
risks = { # type: (loss, bound-coeff, distribution-type, kl factor)
"exact": (None, 1., "dirichlet", 1.),
"MC": (lambda x, y, z: sigmoid_loss(x, y, z, c=cfg.training.sigmoid_c), 1., "dirichlet", 1.),
"Rnd": (lambda x, y, z: rand_loss(x, y, z, n=cfg.training.rand_n), 2., "categorical", cfg.training.rand_n),
"FO": (lambda x, y, z: moment_loss(x, y, z, order=1), 2., "categorical", 1.),
"SO": (lambda x, y, z: moment_loss(x, y, z, order=2), 4., "categorical", 2.),
}
train_errors, test_errors, train_losses, bounds, entropies, times = [], [], [], [], [], []
for i in range(cfg.num_trials):
deterministic(int(cfg.training.seed)+i)
data = Dataset(cfg.dataset.distr, n_train=cfg.dataset.N_train, n_test=cfg.dataset.N_test, noise=cfg.dataset.noise)
if cfg.model.pred == "stumps-uniform":
predictors, M = uniform_decision_stumps(cfg.model.M, 2, data.X_train.min(0), data.X_train.max(0))
elif cfg.model.pred == "stumps-custom":
predictors, M = custom_decision_stumps(torch.zeros((2, 2)), torch.tensor([[1, -1], [1, -1]]))
elif cfg.model.pred == "rf": # random forest
if cfg.model.tree_depth == "None":
cfg.model.tree_depth = None
predictors, M = two_forests(cfg.model.M, 0.5, data.X_train, data.y_train, max_samples=cfg.model.bootstrap, max_depth=cfg.model.tree_depth, binary=data.binary)
else:
raise NotImplementedError("model.pred should be one the following: [stumps-uniform, stumps-custom, rf]")
train_x, train_y, test_x, test_y = torch.from_numpy(data.X_train).float(), torch.from_numpy(data.y_train).float().unsqueeze(1), torch.from_numpy(data.X_test).float(), torch.from_numpy(data.y_test).unsqueeze(1).float()
monitor = MonitorMV(SAVE_DIR)
loss, coeff, distr, kl_factor = risks[cfg.training.risk]
if cfg.model.pred == "rf":
betas = [torch.ones(M) * cfg.model.prior for p in predictors] # prior
# weights proportional to data sizes
model = MultipleMajorityVote(predictors, betas, weights=(0.5, 0.5), mc_draws=cfg.training.MC_draws, distr=distr, kl_factor=kl_factor)
else:
betas = torch.ones(M) * cfg.model.prior # prior
model = MajorityVote(predictors, betas, mc_draws=cfg.training.MC_draws, distr=distr, kl_factor=kl_factor)
# get voter predictions
if cfg.model.pred == "rf":
m = cfg.dataset.N_train // 2 # number of points for learning first prior
# use first m data for learning the second posterior, and the remainder for the first one
train_data = [(train_y[m:], predictors[0](train_x[m:])), (train_y[:m], predictors[1](train_x[:m]))]
# test both posteriors on entire test set
test_data = [(test_y, p(test_x)) for p in predictors]
else:
m = None
train_data = train_y, predictors(train_x)
test_data = test_y, predictors(test_x)
bound = None
if cfg.training.opt_bound:
print(f"Optimize {cfg.bound.type} bound")
bound = lambda n, model, risk: BOUNDS[cfg.bound.type](n, model, risk, cfg.bound.delta, coeff=coeff, monitor=monitor)
optimizer = Adam(model.parameters(), lr=cfg.training.lr)
t1 = time()
train_batch(cfg.dataset.N_train, train_data, model, optimizer, bound=bound, loss=loss, nb_iter=cfg.training.iter, monitor=monitor)
t2 = time()
print(f"{t2-t1}s for {cfg.training.iter} iterations")
test_error = model.risk(test_data)
train_error = model.risk(train_data)
print(f"Test error: {test_error.item()}")
if cfg.training.risk in ["exact", "MC"]:
# evaluate bound with error
b = float(BOUNDS[cfg.bound.type](cfg.dataset.N_train, model, train_error, cfg.bound.delta, coeff=coeff, verbose=True))
else:
# evaluate bound with loss
train_loss = model.risk(train_data, loss)
b = float(BOUNDS[cfg.bound.type](cfg.dataset.N_train, model, train_loss, cfg.bound.delta, coeff=coeff, verbose=True))
train_losses.append(train_loss.item())
train_errors.append(train_error.item())
test_errors.append(test_error.item())
bounds.append(b)
entropies.append(model.entropy().item())
times.append(t2-t1)
monitor.close()
if i == 0: # plot only first trial
plot_2D(data, model, bound=b)
# plt.title(f"{cfg.model.pred}, {cfg.bound.type} bound, M={M}")
plt.savefig(SAVE_DIR / f"{cfg.dataset.distr}.pdf", bbox_inches='tight', transparent=True)
plt.clf()
results = {"train-error": (np.mean(train_errors), np.std(train_errors)),"test-error": (np.mean(test_errors), np.std(test_errors)), cfg.bound.type: (np.mean(bounds), np.std(bounds)), "time": (np.mean(times), np.std(times)), "entropy": (np.mean(entropies), np.std(entropies))}
if cfg.training.risk not in ["exact", "MC"]:
results.update({"train-risk": (np.mean(train_losses), np.std(train_losses))})
np.save(SAVE_DIR / "err-b.npy", results)
if __name__ == "__main__":
main()