-
Notifications
You must be signed in to change notification settings - Fork 0
/
sample.py
81 lines (64 loc) · 2.44 KB
/
sample.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
import os
import numpy as np
import optuna
import wandb
class Objective(object):
def __init__(self, seed):
# Setting the seed to always get the same regression.
np.random.seed(seed)
# Building the data generating process.
self.nobs = 1000
self.epsilon = np.random.uniform(size=(self.nobs, 1))
self.real_beta = 5.0
self.real_alpha = 1.0 / 5.0
self.X1 = np.random.normal(loc=10, scale=3, size=(self.nobs, 1))
self.X2 = np.random.normal(loc=10, scale=3, size=(self.nobs, 1))
self.y = self.X1 * self.real_alpha + \
self.X2 * self.real_beta + \
self.epsilon
def __call__(self, trial):
# Parameters.
trial_alpha = trial.suggest_uniform("alpha", low=-10, high=10)
trial_beta = trial.suggest_uniform("beta", low=-10, high=10)
# Starting WandrB run.
config = {"trial_alpha": trial_alpha,
"trial_beta": trial_beta}
run = wandb.init(project="optuna",
name=f"trial_",
group="sampling",
config=config,
reinit=True)
# Prediction and loss.
y_hat = self.X1 * trial_alpha + self.X2 * trial_beta
mse = ((self.y - y_hat) ** 2).mean()
# WandB logging.
with run:
run.log({"mse": mse}, step=trial.number)
return mse
def main():
# Execute an optimization by using an `Objective` instance.
black_box = Objective(seed=4444)
sampler = optuna.samplers.TPESampler(seed=4444)
study = optuna.create_study(direction="minimize",
sampler=sampler)
study.optimize(black_box,
n_trials=100,
show_progress_bar=True)
print(f"True alpha: {black_box.real_alpha}")
print(f"True beta: {black_box.real_beta}")
print(f"Best params: {study.best_params}")
# Create the summary run.
summary = wandb.init(project="optuna",
name="summary",
job_type="logging")
# Getting the study trials.
trials = study.trials
# WandB summary.
for step, trial in enumerate(trials):
# Logging the loss.
summary.log({"mse": trial.value}, step=step)
# Logging the parameters.
for k, v in trial.params.items():
summary.log({k: v}, step=step)
if __name__ == "__main__":
main()