-
Notifications
You must be signed in to change notification settings - Fork 0
/
optimize.py
78 lines (65 loc) · 2.46 KB
/
optimize.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
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from ax.service.ax_client import AxClient, ObjectiveProperties
from train import SteeringTrainer, SteeringTrainingArgs
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--exp-name", help="the name of the experiment")
args = parser.parse_args()
return args
class MSEWithTau(nn.Module):
def __init__(self, tau):
super().__init__()
self.tau = tau
def forward(self, output, target):
mse = F.mse_loss(output, target, reduction='none')
loss = torch.mean(mse * torch.exp(self.tau * output))
return loss
if __name__ == "__main__":
args = parse_args()
def bo_trial(params):
training_args = SteeringTrainingArgs(
batch_size = 256,
num_epochs = 50,
learning_rate = 0.001,
optimizer = torch.optim.Adam,
loss_fn = MSEWithTau(params['tau']),
results_dir = "results/",
exp_name = f"optimize_{args.exp_name}",
epsilons = [0.1, 0.2, 0.5, 1, 2, 5]
)
trainer = SteeringTrainer(training_args)
trainer.train()
return trainer.evaluate()
ax_client = AxClient()
ax_client.create_experiment(
name="dynamic_tau",
parameters=[
{
"name": "tau",
"type": "range",
"bounds": [0., 2.],
"value_type": "float",
}
],
objectives={"accuracy": ObjectiveProperties(minimize=False)}
)
ax_client.attach_trial(
parameters={"tau": 0.}
)
baseline_parameters = ax_client.get_trial_parameters(trial_index=0)
ax_client.complete_trial(trial_index=0, raw_data=bo_trial(baseline_parameters))
for i in range(3):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=trial_index, raw_data=bo_trial(parameters))
best_parameters, values = ax_client.get_best_parameters()
print(best_parameters, values)
with open(f"results/bo_{args.exp_name}.log", "a") as f:
f.write(best_parameters, values)
ax_client.save_to_json_file(f"results/ax_client_{args.exp_name}.json")
# TODO: model results should be optimizer_1epoch_0.log ... optimize_1epoch_24.log for each trial
# how do i add parameters without optimizing them?
# TODO: somehow getting decent accuracy with tao ~ 30
# perhaps limit it within [0., 2.]