def create_default_exporters(
t2r_model,
export_generator,
compare_fn=create_valid_result_smaller):
"""Creates a list of Exporter to export saved models during evaluation.
Args:
t2r_model: The model to be exported.
export_generator: An export_generator.AbstractExportGenerator.
compare_fn: The function used to deterimne the best model to export.
Returns:
A list containing two exporters, one for numpy and another one for
tf_example interface.
"""
if export_generator is None:
export_generator = default_export_generator.DefaultExportGenerator()
# Create pkl of the input to save alongside the exported models
tmpdir = tempfile.mkdtemp()
in_feature_spec = t2r_model.get_feature_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
in_label_spec = t2r_model.get_label_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
t2r_assets = t2r_pb2.T2RAssets()
t2r_assets.feature_spec.CopyFrom(in_feature_spec.to_proto())
t2r_assets.label_spec.CopyFrom(in_label_spec.to_proto())
t2r_assets_filename = os.path.join(tmpdir,
tensorspec_utils.T2R_ASSETS_FILENAME)
tensorspec_utils.write_t2r_assets_to_file(t2r_assets, t2r_assets_filename)
assets = {tensorspec_utils.T2R_ASSETS_FILENAME: t2r_assets_filename}
export_generator.set_specification_from_model(t2r_model)
exporters = []
exporters.append(
tf.estimator.BestExporter(
name='best_exporter_numpy',
compare_fn=compare_fn(),
serving_input_receiver_fn=export_generator
.create_serving_input_receiver_numpy_fn(),
assets_extra=assets))
exporters.append(
tf.estimator.BestExporter(
name='best_exporter_tf_example',
compare_fn=compare_fn(),
serving_input_receiver_fn=export_generator
.create_serving_input_receiver_tf_example_fn(),
assets_extra=assets))
exporters.append(
tf.estimator.LatestExporter(
name='latest_exporter_numpy',
serving_input_receiver_fn=export_generator
.create_serving_input_receiver_numpy_fn(),
assets_extra=assets))
exporters.append(
tf.estimator.LatestExporter(
name='latest_exporter_tf_example',
serving_input_receiver_fn=export_generator
.create_serving_input_receiver_tf_example_fn(),
assets_extra=assets))
return exporters
def create_hooks(
self,
t2r_model,
estimator,
):
if not self._export_dir and not self._lagged_export_dir:
return []
self._export_generator.set_specification_from_model(t2r_model)
warmup_requests_file = self._export_generator.create_warmup_requests_numpy(
batch_sizes=self._batch_sizes_for_export,
export_dir=estimator.model_dir)
in_feature_spec = t2r_model.get_feature_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
in_label_spec = t2r_model.get_label_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
t2r_assets = t2r_pb2.T2RAssets()
t2r_assets.feature_spec.CopyFrom(in_feature_spec.to_proto())
t2r_assets.label_spec.CopyFrom(in_label_spec.to_proto())
def _export_fn(export_dir, global_step):
"""The actual closure function creating the exported model and assets."""
t2r_assets.global_step = global_step
tmpdir = tempfile.mkdtemp()
t2r_assets_filename = os.path.join(
tmpdir, tensorspec_utils.T2R_ASSETS_FILENAME)
tensorspec_utils.write_t2r_assets_to_file(t2r_assets,
t2r_assets_filename)
res = estimator.export_saved_model(
export_dir_base=export_dir,
serving_input_receiver_fn=self._export_generator.
create_serving_input_receiver_numpy_fn(),
assets_extra={
'tf_serving_warmup_requests': warmup_requests_file,
tensorspec_utils.T2R_ASSETS_FILENAME: t2r_assets_filename
})
return res
return [
contrib_tpu.AsyncCheckpointSaverHook(
save_secs=self._save_secs,
checkpoint_dir=estimator.model_dir,
listeners=[
checkpoint_hooks.LaggedCheckpointListener(
export_fn=_export_fn,
num_versions=self._num_versions,
export_dir=self._export_dir,
lagged_export_dir=self._lagged_export_dir)
])
]
def default_create_export_fn(
t2r_model,
estimator,
export_generator
):
"""Create an export function for a device type.
Args:
t2r_model: A T2RModel instance.
estimator: The estimator used for training.
export_generator: An export generator.
Returns:
A callable function which exports a saved model and returns the path.
"""
in_feature_spec = t2r_model.get_feature_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
in_label_spec = t2r_model.get_label_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
t2r_assets = t2r_pb2.T2RAssets()
t2r_assets.feature_spec.CopyFrom(in_feature_spec.to_proto())
t2r_assets.label_spec.CopyFrom(in_label_spec.to_proto())
def _export_fn(export_dir, global_step):
"""The actual closure function creating the exported model and assets."""
# Create additional assets for the exported models
t2r_assets.global_step = global_step
tmpdir = tempfile.mkdtemp()
t2r_assets_filename = os.path.join(tmpdir,
tensorspec_utils.T2R_ASSETS_FILENAME)
tensorspec_utils.write_t2r_assets_to_file(t2r_assets, t2r_assets_filename)
assets = {
tensorspec_utils.T2R_ASSETS_FILENAME: t2r_assets_filename,
}
return estimator.export_saved_model(
export_dir_base=export_dir,
serving_input_receiver_fn=export_generator
.create_serving_input_receiver_numpy_fn(),
assets_extra=assets)
return _export_fn
def _generate_assets(model, export_dir):
in_feature_spec = model.get_feature_specification_for_packing(
mode=tf.estimator.ModeKeys.PREDICT)
in_label_spec = model.get_label_specification_for_packing(
mode=tf.compat.v1.estimator.ModeKeys.PREDICT)
in_feature_spec = tensorspec_utils.filter_required_flat_tensor_spec(
in_feature_spec)
in_label_spec = tensorspec_utils.filter_required_flat_tensor_spec(
in_label_spec)
t2r_assets = t2r_pb2.T2RAssets()
t2r_assets.feature_spec.CopyFrom(in_feature_spec.to_proto())
t2r_assets.label_spec.CopyFrom(in_label_spec.to_proto())
t2r_assets_dir = os.path.join(export_dir,
tensorspec_utils.EXTRA_ASSETS_DIRECTORY)
tf.io.gfile.makedirs(t2r_assets_dir)
t2r_assets_filename = os.path.join(t2r_assets_dir,
tensorspec_utils.T2R_ASSETS_FILENAME)
tensorspec_utils.write_t2r_assets_to_file(t2r_assets, t2r_assets_filename)
def convert(assets_filepath):
"""Converts existing asset pickle based files to t2r proto based assets."""
t2r_assets = t2r_pb2.T2RAssets()
input_spec_filepath = os.path.join(assets_filepath, 'input_specs.pkl')
if not tf.io.gfile.exists(input_spec_filepath):
raise ValueError('No file exists for {}.'.format(input_spec_filepath))
feature_spec, label_spec = tensorspec_utils.load_input_spec_from_file(
input_spec_filepath)
t2r_assets.feature_spec.CopyFrom(feature_spec.to_proto())
t2r_assets.label_spec.CopyFrom(label_spec.to_proto())
global_step_filepath = os.path.join(assets_filepath, 'global_step.pkl')
if tf.io.gfile.exists(global_step_filepath):
global_step = tensorspec_utils.load_input_spec_from_file(
global_step_filepath)
t2r_assets.global_step = global_step
t2r_assets_filepath = os.path.join(assets_filepath,
tensorspec_utils.T2R_ASSETS_FILENAME)
tensorspec_utils.write_t2r_assets_to_file(t2r_assets, t2r_assets_filepath)
