temp_hyperband.py

import collections
import numpy as np

import ray

ray.init()

def gen_random_hyperparams():
  return (np.random.normal(),)

def get_hyperparameter_configuration(n):
  return [gen_random_hyperparams() for _ in range(n)]

@ray.remote
def run_then_return_val_loss(hyperparam_config, time):
  return time * np.random.uniform() * hyperparam_config[0]

def top_k(configs, losses, k):
  return [configs[i] for i in np.argsort(losses)[:k]]

eta = 3
R = 100

s_max = int(np.log(R) / np.log(eta))
B = (s_max + 1) * R

@ray.remote
def successive_halving(s, n, r, T):
  all_results = collections.defaultdict(lambda: [])
  for i in range(s + 1):
    n_i = int(np.floor(n * eta ** (-i)))
    r_i = r * eta ** i
    L = ray.get([run_then_return_val_loss.remote(t, r_i) for t in T])

    for t, l in zip(T, L):
      all_results[t].append((r_i, l))

    k = n_i // eta
    T = top_k(T, L, k)
  return dict(all_results)

results = []
for s in range(s_max, -1, -1):
  n = int(np.ceil((B / R) * eta ** s / (s + 1)))
  r = R * eta ** (-s)
  T = get_hyperparameter_configuration(n)
  results.append(successive_halving.remote(s, n, r, T))

all_config_results = {}
for result in ray.get(results):
  all_config_results.update(result)

print(all_config_results)

import IPython
IPython.embed()