def test_evaluate(self):
es_size_2 = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=2,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
es_size_4 = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=4,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
metrics = [tf.keras.metrics.MeanSquaredError()]
es_size_2.fit(self.fit_examples, self.fit_label)
es_size_4.fit(self.fit_examples, self.fit_label)
es_2_mse = es_size_2.evaluate(self.fit_examples, self.fit_label, metrics)[0]
es_4_mse = es_size_4.evaluate(self.fit_examples, self.fit_label, metrics)[0]
self.assertLessEqual(es_4_mse, es_2_mse)
def testDoWithMajoritVoting(self):
exec_properties = self._exec_properties.copy()
exec_properties['tuner_fn'] = '%s.%s' % (
tuner_module.tuner_fn.__module__, tuner_module.tuner_fn.__name__)
exec_properties['metalearning_algorithm'] = 'majority_voting'
input_dict = self._input_dict.copy()
ps_type = ps_pb2.Type(
binary_classification=ps_pb2.BinaryClassification(label='class'))
ps = ps_pb2.ProblemStatement(
owner=['nitroml'],
tasks=[ps_pb2.Task(
name='mockdata_1',
type=ps_type,
)])
exec_properties['custom_config'] = json_utils.dumps({
'problem_statement':
text_format.MessageToString(message=ps, as_utf8=True),
})
hps_artifact = artifacts.KCandidateHyperParameters()
hps_artifact.uri = os.path.join(self._testdata_dir,
'MetaLearner.majority_voting',
'hparams_out')
input_dict['warmup_hyperparameters'] = [hps_artifact]
tuner = executor.Executor(self._context)
tuner.Do(input_dict=input_dict,
output_dict=self._output_dict,
exec_properties=exec_properties)
self._verify_output()
def testDoWithTunerFn(self):
self._exec_properties['tuner_fn'] = '%s.%s' % (
tuner_module.tuner_fn.__module__, tuner_module.tuner_fn.__name__)
ps_type = ps_pb2.Type(
binary_classification=ps_pb2.BinaryClassification(label='class'))
ps = ps_pb2.ProblemStatement(
owner=['nitroml'],
tasks=[ps_pb2.Task(
name='mockdata_1',
type=ps_type,
)])
self._exec_properties['custom_config'] = json_utils.dumps({
'problem_statement':
text_format.MessageToString(message=ps, as_utf8=True),
})
tuner = executor.Executor(self._context)
tuner.Do(input_dict=self._input_dict,
output_dict=self._output_dict,
exec_properties=self._exec_properties)
self._verify_output()
def problem_statement(self) -> ps_pb2.ProblemStatement:
"""Returns the ProblemStatement associated with this Task."""
return ps_pb2.ProblemStatement(
owner=['nitroml'],
tasks=[ps_pb2.Task(
name=self.name,
type=self._get_task_type(),
)])
def preprocessing_fn(inputs: Dict[str, Tensor],
custom_config=Dict[str, Any]) -> Dict[str, Tensor]:
"""tf.transform's callback function for preprocessing inputs.
Args:
inputs: map from feature keys to raw not-yet-transformed features.
custom_config: Custom configuration dictionary for passing the task's
ProblemStatement as a text proto, since custom_config must be
JSON-serializable.
Returns:
Map from string feature key to transformed feature operations.
"""
problem_statement = ps_pb2.ProblemStatement()
text_format.Parse(
text=custom_config[BasicPreprocessor.PROBLEM_STATEMENT_KEY],
message=problem_statement)
outputs = {}
for key in [k for k, v in inputs.items() if v.dtype == tf.float32]:
# TODO(weill): Handle case when an int field can actually represents numeric
# rather than categorical values.
task_type = problem_statement.tasks[0].type
if task_type.HasField('one_dimensional_regression') and (
key == task_type.one_dimensional_regression.label):
outputs[key] = inputs[key]
# Skip normalizing regression tasks.
continue
# Preserve this feature as a dense float, setting nan's to the mean.
outputs[_sanitize_feature_name(key)] = tft.scale_to_z_score(
_fill_in_missing(inputs[key]))
for key in [k for k, v in inputs.items() if v.dtype != tf.float32]:
# Build a vocabulary for this feature.
# TODO(weill): Risk here to blow up computation needlessly.
output = tft.compute_and_apply_vocabulary(_fill_in_missing(
inputs[key]),
top_k=None,
num_oov_buckets=1)
# Don't sanitize the label key name.
task_type = problem_statement.tasks[0].type
if task_type.HasField('multi_class_classification') and (
key == task_type.multi_class_classification.label):
outputs[key] = output
continue
if task_type.HasField('binary_classification') and (
key == task_type.binary_classification.label):
outputs[key] = output
continue
# Do sanitize feature key names.
outputs[_sanitize_feature_name(key)] = output
return outputs
def problem_statement(self) -> ps_pb2.ProblemStatement:
"""Returns the ProblemStatement associated with this BenchmarkTask."""
return ps_pb2.ProblemStatement(
owner=['nitroml'],
tasks=[
ps_pb2.Task(
name='Test',
type=ps_pb2.Type(one_dimensional_regression=ps_pb2.
OneDimensionalRegression(label='test')),
)
])
def test_lifecycle(self):
es = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=3,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
test_dir = os.path.join(self.data_path, 'test_examples.tfrecord')
test_examples = np.asarray(
list(tf.data.TFRecordDataset(test_dir).as_numpy_iterator()))
test_examples_tensor = tf.convert_to_tensor(test_examples)
model_predictions = {}
for model_id, path in self.saved_model_paths.items():
reloaded_model = tf.saved_model.load(path)
model_predictions[model_id] = reloaded_model.signatures[
'serving_default'](test_examples_tensor)['output_0'].numpy()
want_weights = {'2': 0.3333333333333333, '4': 0.6666666666666666}
want_prediction = want_weights['2'] * model_predictions['2'] + want_weights[
'4'] * model_predictions['4']
mse = tf.keras.metrics.MeanSquaredError()
mse_scores = []
for pred in model_predictions.values():
mse_scores.append(mse(self.fit_label, pred))
mse.reset_states()
export_dir = os.path.join(
tempfile.mkdtemp(dir=absltest.get_default_test_tmpdir()),
'from_estimator')
es.fit(self.fit_examples, self.fit_label)
ensemble_predictions = es.predict(test_examples)
ensemble_mse = es.evaluate(self.fit_examples, self.fit_label, [mse])[0]
ensemble_path = es.save(export_dir)
reloaded_ensemble = tf.saved_model.load(ensemble_path)
loaded_ensemble_prediction = reloaded_ensemble.signatures[
'serving_default'](input=test_examples_tensor)['output'].numpy()
self.assertEqual(want_weights, es.weights)
self.assertEqual((10, 1), ensemble_predictions.shape)
np.testing.assert_array_almost_equal(want_prediction, ensemble_predictions,
1)
self.assertLessEqual(ensemble_mse, min(mse_scores))
np.testing.assert_array_almost_equal(ensemble_predictions,
loaded_ensemble_prediction, 1)
def problem_statement(self) -> ps_pb2.ProblemStatement:
"""Returns the ProblemStatement associated with this Task."""
# Supervised keys is a two-tuple.
_, target_key = self._dataset_builder.info.supervised_keys
return ps_pb2.ProblemStatement(
owner=['nitroml'],
tasks=[
ps_pb2.Task(
name=self.name,
type=ps_pb2.Type(
binary_classification=ps_pb2.BinaryClassification(
label=target_key)),
)
])
def ensemble_selection(
problem_statement: Parameter[str],
examples: InputArtifact[standard_artifacts.Examples],
evaluation_split_name: Parameter[str],
ensemble_size: Parameter[int],
metric: Parameter[str],
goal: Parameter[str],
model: OutputArtifact[standard_artifacts.Model],
input_model0: InputArtifact[standard_artifacts.Model] = None,
input_model1: InputArtifact[standard_artifacts.Model] = None,
input_model2: InputArtifact[standard_artifacts.Model] = None,
input_model3: InputArtifact[standard_artifacts.Model] = None,
input_model4: InputArtifact[standard_artifacts.Model] = None,
input_model5: InputArtifact[standard_artifacts.Model] = None,
input_model6: InputArtifact[standard_artifacts.Model] = None,
input_model7: InputArtifact[standard_artifacts.Model] = None,
input_model8: InputArtifact[standard_artifacts.Model] = None,
input_model9: InputArtifact[standard_artifacts.Model] = None,
) -> None: # pytype: disable=invalid-annotation,wrong-arg-types
"""Runs the SimpleML trainer as a separate component."""
problem_statement = text_format.Parse(problem_statement,
ps_pb2.ProblemStatement())
input_models = [
input_model0, input_model1, input_model2, input_model3, input_model4,
input_model5, input_model6, input_model7, input_model8, input_model9
]
saved_model_paths = {
str(i): path_utils.serving_model_path(model.uri)
for i, model in enumerate(input_models) if model
}
logging.info('Saved model paths: %s', saved_model_paths)
label_key = _label_key(problem_statement)
es = es_lib.EnsembleSelection(problem_statement=problem_statement,
saved_model_paths=saved_model_paths,
ensemble_size=ensemble_size,
metric=tf.keras.metrics.deserialize(
json.loads(metric)),
goal=goal)
es.fit(*_data_from_examples(examples_path=os.path.join(
examples.uri, evaluation_split_name),
label_key=label_key))
logging.info('Selected ensemble weights: %s', es.weights)
es.save(export_path=os.path.join(path_utils.serving_model_dir(model.uri),
'export', 'serving'))
def test_get_predictions(self):
# TODO(liumich): improve test predictions with the following steps
# - reduce the number of samples to 4-5
# - manually compute MSE after each iteration for each partial ensemble
# - also output the ground truth (labels) so that we can verify
es = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=3,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
want_predictions = {
'0':
np.array([[268520.7], [172055.8], [172840.52], [203374.36],
[629715.5], [160393.], [242507.27], [156286.08],
[262261.7], [221169.3]]),
'1':
np.array([[262822.53], [168104.17], [168874.69], [198855.67],
[617477.56], [156652.53], [237280.08], [152620.],
[256676.81], [216328.45]]),
'2':
np.array([[247936.98], [158487.19], [159214.84], [187528.2],
[582864.9], [147672.52], [223815.3], [143864.3],
[242133.11], [204029.08]]),
'3':
np.array([[268206.75], [171761.81], [172546.36], [203073.86],
[629326.7], [160101.38], [242198.66], [155995.36],
[261948.98], [220865.14]]),
'4':
np.array([[257493.5], [164639.19], [165394.55], [194785.5],
[605169.06], [153412.9], [232453.73], [149459.73],
[251468.73], [211914.42]])
}
predictions = es._get_predictions_dict(self.fit_examples)
for model_id in predictions.keys():
np.testing.assert_array_almost_equal(want_predictions[model_id],
predictions[model_id], 1)
def test_predict_before_fit(self):
es = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=3,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
with self.assertRaisesRegex(
ValueError,
'Weights cannot be empty. Must call `fit` before `predict`.'):
_ = es.predict(self.fit_examples)
def test_calculate_weights(self):
es = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=4,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
ensemble_count = {'model_1': 1, 'model_2': 2, 'model_3': 1}
want_weights = {'model_1': 0.25, 'model_2': 0.5, 'model_3': 0.25}
es._calculate_weights(ensemble_count)
self.assertEqual(want_weights, es.weights)
def test_evaluate_metrics(self):
es = ensemble_selection.EnsembleSelection(
problem_statement=ps_pb2.ProblemStatement(tasks=[
ps_pb2.Task(
type=ps_pb2.Type(
one_dimensional_regression=ps_pb2.OneDimensionalRegression(
label='label')))
]),
saved_model_paths=self.saved_model_paths,
predict_fn=_test_predict_fn,
ensemble_size=3,
metric=tf.keras.metrics.MeanSquaredError(),
goal='minimize')
metrics = [
tf.keras.metrics.MeanSquaredError(),
tf.keras.metrics.MeanAbsoluteError(),
tf.keras.metrics.RootMeanSquaredError()
]
es.fit(self.fit_examples, self.fit_label)
ensemble_metrics = es.evaluate(self.fit_examples, self.fit_label, metrics)
self.assertLen(ensemble_metrics, len(metrics))
def tuner_fn(fn_args: fn_args_utils.FnArgs) -> TunerFnResult:
"""Build the tuner using the KerasTuner API.
Args:
fn_args: Holds args as name/value pairs.
- working_dir: working dir for tuning.
- train_files: List of file paths containing training tf.Example data.
- eval_files: List of file paths containing eval tf.Example data.
- train_steps: number of train steps.
- eval_steps: number of eval steps.
- schema_path: optional schema of the input data.
- transform_graph_path: optional transform graph produced by TFT.
- custom_config: A dict with a single 'problem_statement' entry containing
a text-format serialized ProblemStatement proto which defines the task.
Returns:
A namedtuple contains the following:
- tuner: A BaseTuner that will be used for tuning.
- fit_kwargs: Args to pass to tuner's run_trial function for fitting the
model , e.g., the training and validation dataset. Required
args depend on the above tuner's implementation.
"""
problem_statement = text_format.Parse(
fn_args.custom_config['problem_statement'], ps_pb2.ProblemStatement())
autodata_adapter = kma.KerasModelAdapter(
problem_statement=problem_statement,
transform_graph_dir=fn_args.transform_graph_path)
build_keras_model_fn = functools.partial(_build_keras_model,
autodata_adapter=autodata_adapter)
if 'warmup_hyperparameters' in fn_args.custom_config:
hyperparameters = hp_module.HyperParameters.from_config(
fn_args.custom_config['warmup_hyperparameters'])
else:
hyperparameters = _get_hyperparameters()
tuner_cls = get_tuner_cls_with_callbacks(kerastuner.RandomSearch)
tuner = tuner_cls(build_keras_model_fn,
max_trials=fn_args.custom_config.get('max_trials', 10),
hyperparameters=hyperparameters,
allow_new_entries=False,
objective=autodata_adapter.tuner_objective,
directory=fn_args.working_dir,
project_name=f'{problem_statement.tasks[0].name}_tuning')
# TODO(nikhilmehta): Make batch-size tunable hyperparameter.
train_dataset = autodata_adapter.get_dataset(
file_pattern=fn_args.train_files,
batch_size=128,
num_epochs=None,
shuffle=True)
eval_dataset = autodata_adapter.get_dataset(
file_pattern=fn_args.eval_files,
batch_size=128,
num_epochs=1,
shuffle=False)
return TunerFnResult(tuner=tuner,
fit_kwargs={
'x': train_dataset,
'validation_data': eval_dataset,
'steps_per_epoch': fn_args.train_steps,
})
def run_fn(fn_args: trainer_executor.TrainerFnArgs):
"""Train a DNN Keras Model based on given args.
Args:
fn_args: Holds args used to train the model as name/value pairs.
- train_files: A list of uris for train files.
- transform_output: An optional single uri for transform graph produced by
TFT. Will be None if not specified.
- serving_model_dir: A single uri for the output directory of the serving
model.
- eval_model_dir: A single uri for the output directory of the eval model.
Note that this is for estimator only, Keras doesn't require it for TFMA.
- eval_files: A list of uris for eval files.
- schema_file: A single uri for schema file.
- train_steps: Number of train steps.
- eval_steps: Number of eval steps.
- base_model: Base model that will be used for this training job.
- hyperparameters: An optional kerastuner.HyperParameters config.
- custom_config: A dict with a single 'problem_statement' entry containing
a text-format serialized ProblemStatement proto which defines the task.
"""
# Use EstimatorAdapter here because we will wrap the Keras Model into an
# Estimator for training and export.
autodata_adapter = ea.EstimatorAdapter(
problem_statement=text_format.Parse(
fn_args.custom_config['problem_statement'],
ps_pb2.ProblemStatement()),
transform_graph_dir=fn_args.transform_output)
if fn_args.hyperparameters:
hparams = kerastuner.HyperParameters.from_config(
fn_args.hyperparameters)
else:
hparams = _get_hyperparameters()
logging.info('HyperParameters for training: %s', hparams.get_config())
# Use KerasAdapter here because we create need it to create the Keras Model.
keras_autodata_adapter = kma.KerasModelAdapter(
problem_statement=text_format.Parse(
fn_args.custom_config['problem_statement'],
ps_pb2.ProblemStatement()),
transform_graph_dir=fn_args.transform_output)
model = _build_keras_model(hparams,
keras_autodata_adapter,
sequence_length=fn_args.custom_config.get(
'sequence_length', None))
train_spec = tf.estimator.TrainSpec(input_fn=autodata_adapter.get_input_fn(
file_pattern=fn_args.train_files,
batch_size=64,
num_epochs=None,
shuffle=True),
max_steps=fn_args.train_steps)
serving_receiver_fn = autodata_adapter.get_serving_input_receiver_fn()
exporters = [
tf.estimator.FinalExporter('serving_model_dir', serving_receiver_fn),
]
eval_spec = tf.estimator.EvalSpec(
input_fn=autodata_adapter.get_input_fn(file_pattern=fn_args.eval_files,
batch_size=64,
num_epochs=1,
shuffle=False),
steps=fn_args.eval_steps,
exporters=exporters,
# Since eval runs in parallel, we can begin evaluation as soon as new
# checkpoints are written.
start_delay_secs=1,
throttle_secs=5)
run_config = tf.estimator.RunConfig(model_dir=fn_args.serving_model_dir,
save_checkpoints_steps=999,
keep_checkpoint_max=3)
estimator = tf.keras.estimator.model_to_estimator(model, config=run_config)
logging.info('Training model...')
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
logging.info('Training complete.')
# Export an eval savedmodel for TFMA. If distributed training, it must only
# be written by the chief worker, as would be done for serving savedmodel.
if run_config.is_chief:
logging.info('Exporting eval_savedmodel for TFMA.')
tfma.export.export_eval_savedmodel(
estimator=estimator,
export_dir_base=fn_args.eval_model_dir,
eval_input_receiver_fn=autodata_adapter.get_eval_input_receiver_fn(
))
logging.info('Exported eval_savedmodel to %s.', fn_args.eval_model_dir)
else:
logging.info('eval_savedmodel export for TFMA is skipped because '
'this is not the chief worker.')
def encode(
self,
component_spec: Optional[tfx_types.ComponentSpec] = None
) -> message.Message:
# Return an arbitrary proto.
return ps_pb2.ProblemStatement()
def run_fn(fn_args: trainer_executor.TrainerFnArgs):
"""Train a DNNEstimator based on given args.
Args:
fn_args: Holds args used to train the model as name/value pairs.
- train_files: A list of uris for train files.
- transform_output: An optional single uri for transform graph produced by
TFT. Will be None if not specified.
- serving_model_dir: A single uri for the output directory of the serving
model.
- eval_model_dir: A single uri for the output directory of the eval model.
Note that this is estimator only, Keras doesn't require it for TFMA.
- eval_files: A list of uris for eval files.
- schema_file: A single uri for schema file.
- train_steps: Number of train steps.
- eval_steps: Number of eval steps.
- base_model: Base model that will be used for this training job.
- hyperparameters: An optional kerastuner.HyperParameters config.
- custom_config: A dict with a single 'problem_statement' entry containing
a text-format serialized ProblemStatement proto which defines the task.
"""
sequence_length = fn_args.custom_config.get('sequence_length', None)
if sequence_length:
raise ValueError('Sequential prediction tasks are not supported. '
'Set `use_keras=True` in AutoTrainer instead.')
autodata_adapter = estimator_adapter.EstimatorAdapter(
problem_statement=text_format.Parse(
fn_args.custom_config['problem_statement'],
ps_pb2.ProblemStatement()),
transform_graph_dir=fn_args.transform_output)
run_config = tf.estimator.RunConfig(model_dir=fn_args.serving_model_dir,
save_checkpoints_steps=999,
keep_checkpoint_max=3)
estimator = tf.estimator.DNNEstimator(
head=autodata_adapter.head,
feature_columns=autodata_adapter.get_dense_feature_columns(),
hidden_units=[128, 128],
config=run_config)
train_spec = tf.estimator.TrainSpec(input_fn=autodata_adapter.get_input_fn(
file_pattern=fn_args.train_files,
batch_size=64,
num_epochs=None,
shuffle=True),
max_steps=fn_args.train_steps)
serving_receiver_fn = autodata_adapter.get_serving_input_receiver_fn()
exporters = [
tf.estimator.FinalExporter('serving_model_dir', serving_receiver_fn),
]
eval_spec = tf.estimator.EvalSpec(
input_fn=autodata_adapter.get_input_fn(file_pattern=fn_args.eval_files,
batch_size=64,
num_epochs=1,
shuffle=False),
steps=fn_args.eval_steps,
exporters=exporters,
# Since eval runs in parallel, we can begin evaluation as soon as new
# checkpoints are written.
start_delay_secs=1,
throttle_secs=5)
# Train/Tune the model
logging.info('Training model...')
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
logging.info('Training complete.')
# Export an eval savedmodel for TFMA. If distributed training, it must only
# be written by the chief worker, as would be done for serving savedmodel.
if run_config.is_chief:
logging.info('Exporting eval_savedmodel for TFMA.')
tfma.export.export_eval_savedmodel(
estimator=estimator,
export_dir_base=fn_args.eval_model_dir,
eval_input_receiver_fn=autodata_adapter.get_eval_input_receiver_fn(
))
logging.info('Exported eval_savedmodel to %s.', fn_args.eval_model_dir)
else:
logging.info('eval_savedmodel export for TFMA is skipped because '
'this is not the chief worker.')
