def test_fixed_knobs(self, budget):
knob_config = {
'fixed': FixedKnob('fixed'),
'fixed2': FixedKnob(2),
'early_stop': PolicyKnob('EARLY_STOP')
}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, FixedAdvisor)
def test_standard_knobs(self, budget):
knob_config = {
'int': IntegerKnob(2, 128),
'float': FloatKnob(1e-5, 1e-1, is_exp=True),
'cat': CategoricalKnob([16, 32, 64, 128]),
}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, BayesOptAdvisor)
def test_arch_knobs_with_standard_knobs(self, budget, arch_knob):
knob_config = {
'arch': arch_knob,
'float': FloatKnob(1e-5, 1e-1, is_exp=True),
'cat': CategoricalKnob([16, 32, 64, 128]),
}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, RandomAdvisor)
def test_standard_knobs_with_params_sharing(self, budget):
knob_config = {
'int': IntegerKnob(2, 128),
'float': FloatKnob(1e-5, 1e-1, is_exp=True),
'cat': CategoricalKnob([16, 32, 64, 128]),
'share_params': PolicyKnob('SHARE_PARAMS')
}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, BayesOptWithParamSharingAdvisor)
def test_standard_knobs_with_early_stop(self, budget):
knob_config = {
'int': IntegerKnob(2, 128),
'float': FloatKnob(1e-5, 1e-1, is_exp=True),
'cat': CategoricalKnob([16, 32, 64, 128]),
'early_stop': PolicyKnob('EARLY_STOP')
}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, BayesOptAdvisor)
def test_arch_knob_without_enas_policies(self, day_budget, arch_knob):
knob_config = {
'arch': arch_knob,
'downscale': PolicyKnob('DOWNSCALE'),
'early_stop': PolicyKnob('EARLY_STOP'),
'skip_train': PolicyKnob('SKIP_TRAIN')
# No quick eval
}
advisor = make_advisor(knob_config, day_budget)
assert isinstance(advisor, RandomAdvisor)
def _make_advisor(self):
clazz = self._monitor.model_class
budget = self._monitor.budget
# Retrieve knob config
knob_config = clazz.get_knob_config()
advisor = make_advisor(knob_config, budget)
logger.info(f'Using advisor "{type(advisor).__name__}"...')
return advisor
def test_arch_knob_without_budget(self, budget, arch_knob):
knob_config = {
'arch': arch_knob,
'fixed': FixedKnob('fixed'),
'downscale': PolicyKnob('DOWNSCALE'),
'early_stop': PolicyKnob('EARLY_STOP'),
'skip_train': PolicyKnob('SKIP_TRAIN'),
'quick_eval': PolicyKnob('QUICK_EVAL')
}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, RandomAdvisor)
def test_arch_knob_with_enas_policies(self, day_budget, arch_knob):
knob_config = {
'arch': arch_knob,
'fixed': FixedKnob('fixed'),
'share_params': PolicyKnob('SHARE_PARAMS'),
'downscale': PolicyKnob('DOWNSCALE'),
'early_stop': PolicyKnob('EARLY_STOP'),
'skip_train': PolicyKnob('SKIP_TRAIN'),
'quick_eval': PolicyKnob('QUICK_EVAL')
}
advisor = make_advisor(knob_config, day_budget)
assert isinstance(advisor, EnasAdvisor)
def tune_model(
py_model_class: Type[BaseModel],
train_dataset_path: str,
val_dataset_path: str,
test_dataset_path: str = None,
budget: Budget = None,
train_args: Dict[str, any] = None) -> (Dict[str, Any], float, Params):
worker_id = 'local'
# Note start time
start_time = time.time()
# Retrieve config of model
_print_header('Checking model configuration...')
knob_config = py_model_class.get_knob_config()
_check_knob_config(knob_config)
# Read knob values from CLI args
_print_header('Starting trials...')
knobs_from_args = _maybe_read_knobs_from_args(knob_config)
# Read budget options from CLI args
budget_from_args = _maybe_read_budget_from_args()
budget = {**(budget or {}), **budget_from_args}
inform_user(f'Using budget {budget}...')
# Make advisor
advisor = make_advisor(knob_config, budget)
inform_user(f'Using advisor "{type(advisor).__name__}"...')
# Create caches & stores
param_store: ParamStore = FileParamStore()
param_cache: ParamCache = ParamCache()
train_cache: TrainCache = TrainCache()
# Variables to track over trials
best_model_score = -1
best_trial_no = 0
best_model_test_score = None
best_proposal = None
best_store_params_id = None
# Train worker tells advisor that it is free
train_cache.add_worker(worker_id)
# Until there's no more proposals, keep conducting trials
trial_no = 0
while True:
trial_no += 1
# Advisor checks free workers
worker_ids = train_cache.get_workers()
assert worker_id in worker_ids
# Advisor checks worker doesn't already have a proposal
proposal = train_cache.get_proposal(worker_id)
assert proposal is None
# Advisor sends a proposal to worker
# Overriding knobs from args
proposal = advisor.propose(worker_id, trial_no)
if proposal is None:
print('No more proposals from advisor - to stop training')
break
proposal.knobs = {**proposal.knobs, **knobs_from_args}
train_cache.create_proposal(worker_id, proposal)
# Worker receives proposal
proposal = train_cache.get_proposal(worker_id)
assert proposal is not None
# Worker starts trial
_print_header(f'Trial #{trial_no}')
print('Proposal from advisor:', proposal)
# Worker loads model
model_inst = py_model_class(**proposal.knobs)
# Worker pulls shared params
shared_params = _pull_shared_params(proposal, param_cache)
# Worker trains model
print('Training model...')
model_inst.train(train_dataset_path,
shared_params=shared_params,
**(train_args or {}))
# Worker evaluates model
result = _evaluate_model(model_inst, proposal, val_dataset_path)
# Worker caches/saves model parameters
store_params_id = _save_model(model_inst, proposal, result,
param_cache, param_store)
# Update best saved model
if result.score is not None and store_params_id is not None and result.score > best_model_score:
inform_user(
'Best saved model so far! Beats previous best of score {}!'.
format(best_model_score))
best_store_params_id = store_params_id
best_proposal = proposal
best_model_score = result.score
best_trial_no = trial_no
# Test best model, if test dataset provided
if test_dataset_path is not None:
print('Evaluating new best model on test dataset...')
best_model_test_score = model_inst.evaluate(test_dataset_path)
inform_user(
'Score on test dataset: {}'.format(best_model_test_score))
# Worker sends result to advisor
print('Giving feedback to advisor...')
train_cache.create_result(worker_id, result)
train_cache.delete_proposal(worker_id)
# Advisor receives result
# Advisor ingests feedback
result = train_cache.take_result(worker_id)
assert result is not None
advisor.feedback(worker_id, result)
# Destroy model
model_inst.destroy()
# Train worker tells advisor that it is no longer free
train_cache.delete_worker(worker_id)
# Declare best model
if best_proposal is not None:
inform_user('Best trial #{} has knobs {} with score of {}'.format(
best_trial_no, best_proposal.knobs, best_model_score))
if best_model_test_score is not None:
inform_user(
'...with test score of {}'.format(best_model_test_score))
# Load params for best model
best_params = None
if best_store_params_id is not None:
best_params = param_store.load(best_store_params_id)
# Teardown model class
print('Running model class teardown...')
py_model_class.teardown()
# Print duration
duration = time.time() - start_time
print('Tuning took a total of {}s'.format(duration))
return (best_proposal, best_model_test_score, best_params)
def validate_model_class(model_class,
train_dataset_uri,
test_dataset_uri,
task,
queries=[],
knobs=None):
'''
Validates whether a model class is properly defined.
The model instance's methods will be called in an order similar to that in Rafiki.
:param str train_dataset_uri: URI of the train dataset for testing the training of model
:param str test_dataset_uri: URI of the test dataset for testing the evaluating of model
:param str task: Task type of model
:param list[any] queries: List of queries for testing predictions with the trained model
:param knobs: Knobs to train the model with. If not specified, knobs from an advisor will be used
:type knobs: dict[str, any]
:returns: The trained model
'''
print('Testing instantiation of model...')
model_inst = model_class()
print('Testing getting of model\'s knob config...')
knob_config = model_inst.get_knob_config()
if not isinstance(knob_config, dict):
raise InvalidModelClassException(
'`get_knob_config()` should return a dict[str, any]')
if 'knobs' not in knob_config:
raise InvalidModelClassException(
'`knob_config` should have a \'knobs\' key')
advisor = make_advisor(knob_config)
if knobs is None:
knobs = advisor.propose()
print('Testing initialization of model...')
print('Using knobs: {}'.format(knobs))
model_inst.init(knobs)
print('Testing training of model...')
model_inst.train(train_dataset_uri, task)
print('Testing evaluation of model...')
score = model_inst.evaluate(test_dataset_uri, task)
if not isinstance(score, float):
raise InvalidModelClassException('`evaluate()` should return a float!')
print('Score: {}'.format(score))
print('Testing dumping of parameters of model...')
parameters = model_inst.dump_parameters()
if not isinstance(parameters, dict):
raise InvalidModelClassException(
'`dump_parameters()` should return a dict[str, any]')
try:
json.dumps(parameters)
except Exception:
traceback.print_stack()
raise InvalidModelClassException(
'`parameters` should be JSON serializable')
print('Testing destroying of model...')
model_inst.destroy()
print('Testing loading of parameters of model...')
model_inst = model_class()
model_inst.init(knobs)
model_inst.load_parameters(parameters)
print('Testing retrieving of predict label mapping...')
predict_label_mapping = model_inst.get_predict_label_mapping()
if not isinstance(predict_label_mapping, dict):
raise InvalidModelClassException(
'`get_predict_label_mapping()` should return a dict[int, str]')
print('Testing predictions with model...')
print('Using queries: {}'.format(queries))
predictions = model_inst.predict(queries)
predictions = ensemble_predictions([predictions], [predict_label_mapping],
task)
try:
for prediction in predictions:
prediction = parse_model_prediction(prediction)
json.dumps(prediction)
except Exception:
traceback.print_stack()
raise InvalidModelClassException(
'Each `prediction` should be JSON serializable')
print('Predictions: {}'.format(predictions))
print('The model definition is valid!')
return model_inst
def test_arch_knobs(self, budget, arch_knob):
knob_config = {'arch': arch_knob, 'arch2': arch_knob}
advisor = make_advisor(knob_config, budget)
assert isinstance(advisor, RandomAdvisor)
