import copy
import pickle
from polaris.params import Domain
from polaris.rabbitmq import JobClient
[docs]class Polaris(object):
"""
A client for the Polaris.
"""
def __init__(
self, fn, bounds, algo, trials,
max_evals=10, run_once=False, exp_key=None,
logger=None, debug=False, args=None):
"""
Parameters
----------
fn : callable
fnction which Polaris will evaluate
bounds : array of params
search range of hyperparameters
algo : str
search algorithms (random, tpe, bayesian)
trials : polaris.trials
trials object which store all parameters of trials
max_evals : int
the maximum number of evaluations
exp_key : str
the expriment key to distinguish each experiments
logger : Logger
user logger object
debug : bool
make Polaris debug mode
args : tuple
args to pass fn
"""
self.fn = fn
self.bounds = bounds
self.algo = algo
self.trials = trials
self.max_evals = max_evals
self.run_once = run_once
self.debug = debug
self.args = args
if exp_key is None:
self.exp_key = fn.__name__
else:
self.exp_key = exp_key
self.exp_info = {
'exp_key': exp_key,
'eval_count': len(trials) + 1,
}
if logger is None:
import logging
self.logger = logging.getLogger(__name__)
if self.debug:
self.logger.setLevel(logging.DEBUG)
else:
self.logger.setLevel(logging.INFO)
stream = logging.StreamHandler()
formatter = logging.Formatter(
'%(asctime)s:%(lineno)d:%(levelname)s:%(message)s')
stream.setFormatter(formatter)
self.logger.addHandler(stream)
else:
self.logger = logger
self.domain = Domain(self.bounds, algo=self.algo)
[docs] def run(self, use_mpi=False):
"""
Run an experiment sequentially with one process up to max_evals count.
The expriment will early stop, if params become None
Parameters
----------
use_mpi : boolean
boolean for using mpi or not
"""
if use_mpi:
self.run_with_mpi()
else:
self.run_single()
def run_single(self):
self.logger.info('Start searching...')
for eval_count in range(self.exp_info['eval_count'], self.max_evals+1):
params = self.domain.predict(self.trials)[0]
fn_params = copy.copy(params)
if self.args is None:
exp_result = self.fn(fn_params, self.exp_info)
else:
exp_result = self.fn(fn_params, self.exp_info, *self.args)
self.trials.add(exp_result, params, self.exp_info)
self.exp_info['eval_count'] += 1
self.logger.debug(fn_params)
with open(f'{self.exp_key}.p', mode='wb') as f:
pickle.dump(self.trials, f)
if self.run_once:
break
return self.trials.best_params
def run_with_mpi(self):
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
for eval_count in range(self.exp_info['eval_count'], self.max_evals+1):
if rank == 0:
params = self.domain.predict(self.trials)[0]
fn_params = copy.copy(params)
else:
fn_params = None
fn_params = comm.bcast(fn_params, root=0)
if self.args is None:
exp_result = self.fn(fn_params, self.exp_info)
else:
exp_result = self.fn(fn_params, self.exp_info, *self.args)
if rank == 0:
self.trials.add(exp_result, params, self.exp_info)
self.exp_info['eval_count'] += 1
self.logger.debug(fn_params)
with open(f'{self.exp_key}.p', mode='wb') as f:
pickle.dump(self.trials, f)
if self.run_once:
break
if rank == 0:
return self.trials.best_params
[docs] def run_parallel(self):
"""
Run an experiment in parallel.
Polaris use RabbitMQ to pass arguments from client to worker.
You need to start at least one worker to start an experiment.
"""
job_client = JobClient(self)
job_client.start()
return self.trials.best_params