PySR / hyperparamopt.py
MilesCranmer's picture
Rename hyperopt to fix import issue
9e02588
raw
history blame
5.27 kB
"""Start a hyperoptimization from a single node"""
import sys
import numpy as np
import pickle as pkl
import hyperopt
from hyperopt import hp, fmin, tpe, Trials
import signal
import eureqa
#Change the following code to your file
################################################################################
TRIALS_FOLDER = 'trials'
NUMBER_TRIALS_PER_RUN = 1
def run_trial(args):
"""Evaluate the model loss using the hyperparams in args
:args: A dictionary containing all hyperparameters
:returns: Dict with status and loss from cross-validation
"""
for key in 'niterations npop ncyclesperiteration topn'.split(' '):
args[key] = int(args[key])
def handler(signum, frame):
raise ValueError("Takes too long")
signal.signal(signal.SIGALRM, handler)
maxTime = 60
ntrials = 1 #3
equation_file=f'hall_of_fame_{np.random.rand():f}.csv'
signal.alarm(maxTime)
try:
trials = []
for i in range(1, 2):
trials.append([np.min(eureqa.eureqa(
test=f"simple{i}",
threads=20,
binary_operators=["plus", "mult", "pow", "div"],
unary_operators=["cos", "exp", "sin", "log"],
equation_file=equation_file,
**args)['MSE']) for _ in range(ntrials)])
signal.alarm(0)
except ValueError:
return {
'status': 'ok', # or 'fail' if nan loss
'loss': np.inf
}
loss = np.average(trials)
return {
'status': 'ok', # or 'fail' if nan loss
'loss': loss
}
space = {
'niterations': hp.qloguniform('niterations', np.log(10), 0.5, 1),
'npop': hp.qloguniform('npop', np.log(100), 0.5, 1),
'ncyclesperiteration': hp.qloguniform('ncyclesperiteration', np.log(5000), 0.5, 1),
'topn': hp.quniform('topn', 1, 30, 1),
'annealing': hp.choice('annealing', [False, True]),
'alpha': hp.lognormal('alpha', np.log(10.0), 0.5),
'parsimony': hp.lognormal('parsimony', np.log(1e-3), 0.5),
'fractionReplacedHof': hp.lognormal('fractionReplacedHof', np.log(0.1), 0.5),
'fractionReplaced': hp.lognormal('fractionReplaced', np.log(0.1), 0.5),
'weightMutateConstant': hp.lognormal('weightMutateConstant', np.log(4.0), 0.5),
'weightMutateOperator': hp.lognormal('weightMutateOperator', np.log(0.5), 0.5),
'weightAddNode': hp.lognormal('weightAddNode', np.log(0.5), 0.5),
'weightDeleteNode': hp.lognormal('weightDeleteNode', np.log(0.5), 0.5),
'weightSimplify': hp.lognormal('weightSimplify', np.log(0.05), 0.5),
'weightRandomize': hp.lognormal('weightRandomize', np.log(0.25), 0.5),
'weightDoNothing': hp.lognormal('weightDoNothing', np.log(1.0), 0.5),
}
################################################################################
def merge_trials(trials1, trials2_slice):
"""Merge two hyperopt trials objects
:trials1: The primary trials object
:trials2_slice: A slice of the trials object to be merged,
obtained with, e.g., trials2.trials[:10]
:returns: The merged trials object
"""
max_tid = 0
if len(trials1.trials) > 0:
max_tid = max([trial['tid'] for trial in trials1.trials])
for trial in trials2_slice:
tid = trial['tid'] + max_tid + 1
hyperopt_trial = Trials().new_trial_docs(
tids=[None],
specs=[None],
results=[None],
miscs=[None])
hyperopt_trial[0] = trial
hyperopt_trial[0]['tid'] = tid
hyperopt_trial[0]['misc']['tid'] = tid
for key in hyperopt_trial[0]['misc']['idxs'].keys():
hyperopt_trial[0]['misc']['idxs'][key] = [tid]
trials1.insert_trial_docs(hyperopt_trial)
trials1.refresh()
return trials1
loaded_fnames = []
trials = None
# Run new hyperparameter trials until killed
while True:
np.random.seed()
# Load up all runs:
import glob
path = TRIALS_FOLDER + '/*.pkl'
for fname in glob.glob(path):
if fname in loaded_fnames:
continue
trials_obj = pkl.load(open(fname, 'rb'))
n_trials = trials_obj['n']
trials_obj = trials_obj['trials']
if len(loaded_fnames) == 0:
trials = trials_obj
else:
print("Merging trials")
trials = merge_trials(trials, trials_obj.trials[-n_trials:])
loaded_fnames.append(fname)
print("Loaded trials", len(loaded_fnames))
if len(loaded_fnames) == 0:
trials = Trials()
n = NUMBER_TRIALS_PER_RUN
try:
best = fmin(run_trial,
space=space,
algo=tpe.suggest,
max_evals=n + len(trials.trials),
trials=trials,
verbose=1,
rstate=np.random.RandomState(np.random.randint(1,10**6))
)
except hyperopt.exceptions.AllTrialsFailed:
continue
print('current best', best)
hyperopt_trial = Trials()
# Merge with empty trials dataset:
save_trials = merge_trials(hyperopt_trial, trials.trials[-n:])
new_fname = TRIALS_FOLDER + '/' + str(np.random.randint(0, sys.maxsize)) + '.pkl'
pkl.dump({'trials': save_trials, 'n': n}, open(new_fname, 'wb'))
loaded_fnames.append(new_fname)