def start_bundle(self) -> None:
# There's no initialisation method for CombineFns.
# See BEAM-3736: Add SetUp() and TearDown() for CombineFns.
# Default to eval_saved_model dofn to preserve legacy assumption
# of eval_saved_model.
# TODO(ihchen): Update all callers and make this an error condition to not
# have construct_fn specified.
if not self._counters_incremented:
self._counters_incremented = True
counter_util.update_beam_counters(
self._eval_shared_model.add_metrics_callbacks)
if self._eval_shared_model.construct_fn is None:
construct_fn = dofn.make_construct_fn(
eval_saved_model_path=self._eval_shared_model.model_path,
add_metrics_callbacks=self._eval_shared_model.
add_metrics_callbacks,
include_default_metrics=True,
additional_fetches=None)
self._eval_saved_model = (
self._eval_shared_model.shared_handle.acquire(
construct_fn(self._model_load_seconds)))
else:
self._eval_saved_model = (
self._eval_shared_model.shared_handle.acquire(
self._eval_shared_model.construct_fn(
self._model_load_seconds)))
def _start_bundle(self):
# There's no initialisation method for CombineFns.
# See BEAM-3736: Add SetUp() and TearDown() for CombineFns.
self._eval_saved_model = self._shared_handle.acquire(
dofn.make_construct_fn(self._eval_saved_model_path,
self._add_metrics_callbacks,
self._model_load_seconds))
def createTestEvalSharedModel(self,
eval_saved_model_path,
add_metrics_callbacks=None,
include_default_metrics=True,
example_weight_key=None):
return types.EvalSharedModel(
model_path=eval_saved_model_path,
add_metrics_callbacks=add_metrics_callbacks,
example_weight_key=example_weight_key,
construct_fn=dofn.make_construct_fn(eval_saved_model_path,
add_metrics_callbacks,
include_default_metrics))
def default_eval_shared_model(
eval_saved_model_path: Text,
add_metrics_callbacks: Optional[List[
types.AddMetricsCallbackType]] = None,
include_default_metrics: Optional[bool] = True,
example_weight_key: Optional[Text] = None,
additional_fetches: Optional[List[Text]] = None
) -> types.EvalSharedModel:
"""Returns default EvalSharedModel.
Args:
eval_saved_model_path: Path to EvalSavedModel.
add_metrics_callbacks: Optional list of callbacks for adding additional
metrics to the graph (see EvalSharedModel for more information on how to
configure additional metrics). Metrics for example counts and example
weight will be added automatically.
include_default_metrics: True to include the default metrics that are part
of the saved model graph during evaluation.
example_weight_key: Deprecated.
additional_fetches: Prefixes of additional tensors stored in
signature_def.inputs that should be fetched at prediction time. The
"features" and "labels" tensors are handled automatically and should not
be included.
"""
# Always compute example weight and example count.
# PyType doesn't know about the magic exports we do in post_export_metrics.
# Additionally, the lines seem to get reordered in compilation, so we can't
# just put the disable-attr on the add_metrics_callbacks lines.
# pytype: disable=module-attr
if not add_metrics_callbacks:
add_metrics_callbacks = []
example_count_callback = post_export_metrics.example_count()
add_metrics_callbacks.append(example_count_callback)
# TODO(b/126924645): Remove
if example_weight_key:
example_weight_callback = post_export_metrics.example_weight(
example_weight_key)
add_metrics_callbacks.append(example_weight_callback)
# pytype: enable=module-attr
return types.EvalSharedModel(
model_path=eval_saved_model_path,
add_metrics_callbacks=add_metrics_callbacks,
include_default_metrics=include_default_metrics,
example_weight_key=example_weight_key,
additional_fetches=additional_fetches,
construct_fn=dofn.make_construct_fn(
eval_saved_model_path,
add_metrics_callbacks,
include_default_metrics,
additional_fetches=additional_fetches))
def createTestEvalSharedModel(
self,
eval_saved_model_path: Text,
add_metrics_callbacks: Optional[List[
types.AddMetricsCallbackType]] = None,
include_default_metrics: Optional[bool] = True,
example_weight_key: Optional[Text] = None,
additional_fetches: Optional[List[Text]] = None
) -> types.EvalSharedModel:
return types.EvalSharedModel(
model_path=eval_saved_model_path,
add_metrics_callbacks=add_metrics_callbacks,
example_weight_key=example_weight_key,
construct_fn=dofn.make_construct_fn(eval_saved_model_path,
add_metrics_callbacks,
include_default_metrics,
additional_fetches))
def start_bundle(self):
# There's no initialisation method for CombineFns.
# See BEAM-3736: Add SetUp() and TearDown() for CombineFns.
# Default to eval_saved_model dofn to preserve legacy assumption
# of eval_saved_model.
if self._eval_shared_model.construct_fn is None:
construct_fn = dofn.make_construct_fn(
eval_saved_model_path=self._eval_shared_model.model_path,
add_metrics_callbacks=self._eval_shared_model.
add_metrics_callbacks,
include_default_metrics=True)
self._eval_saved_model = (
self._eval_shared_model.shared_handle.acquire(
construct_fn(self._model_load_seconds)))
else:
self._eval_saved_model = (
self._eval_shared_model.shared_handle.acquire(
self._eval_shared_model.construct_fn(
self._model_load_seconds)))
def _start_bundle(self) -> None:
# There's no initialisation method for CombineFns.
# See BEAM-3736: Add SetUp() and TearDown() for CombineFns.
# Default to eval_saved_model dofn to preserve legacy assumption
# of eval_saved_model.
#
# TODO(b/133761055): Update all callers and make this an error condition to
# not have construct_fn specified.
if self._eval_shared_model.construct_fn is None:
construct_fn = dofn.make_construct_fn(
self._eval_shared_model.model_path,
self._eval_shared_model.add_metrics_callbacks,
self._eval_shared_model.include_default_metrics,
self._eval_shared_model.additional_fetches)
else:
construct_fn = self._eval_shared_model.construct_fn
self._eval_metrics_graph = (
self._eval_shared_model.shared_handle.acquire(
construct_fn(self._model_load_seconds)))
def default_eval_shared_model(eval_saved_model_path,
add_metrics_callbacks=None,
include_default_metrics=True,
example_weight_key=None):
"""Returns default EvalSharedModel.
Args:
eval_saved_model_path: Path to EvalSavedModel.
add_metrics_callbacks: Optional list of callbacks for adding additional
metrics to the graph (see EvalSharedModel for more information on how to
configure additional metrics). Metrics for example counts and example
weight will be added automatically.
include_default_metrics: True to include the default metrics that are part
of the saved model graph during evaluation.
example_weight_key: The key of the example weight column. If None, weight
will be 1 for each example.
"""
# Always compute example weight and example count.
# PyType doesn't know about the magic exports we do in post_export_metrics.
# Additionally, the lines seem to get reordered in compilation, so we can't
# just put the disable-attr on the add_metrics_callbacks lines.
# pytype: disable=module-attr
if not add_metrics_callbacks:
add_metrics_callbacks = []
example_count_callback = post_export_metrics.example_count()
add_metrics_callbacks.append(example_count_callback)
if example_weight_key:
example_weight_callback = post_export_metrics.example_weight(
example_weight_key)
add_metrics_callbacks.append(example_weight_callback)
# pytype: enable=module-attr
return types.EvalSharedModel(
model_path=eval_saved_model_path,
add_metrics_callbacks=add_metrics_callbacks,
include_default_metrics=include_default_metrics,
example_weight_key=example_weight_key,
construct_fn=dofn.make_construct_fn(eval_saved_model_path,
add_metrics_callbacks,
include_default_metrics))
def default_eval_shared_model(
eval_saved_model_path: Text,
add_metrics_callbacks: Optional[List[types.AddMetricsCallbackType]] = None,
include_default_metrics: Optional[bool] = True,
example_weight_key: Optional[Union[Text, Dict[Text, Text]]] = None,
additional_fetches: Optional[List[Text]] = None,
blacklist_feature_fetches: Optional[List[Text]] = None
) -> types.EvalSharedModel:
"""Returns default EvalSharedModel.
Args:
eval_saved_model_path: Path to EvalSavedModel.
add_metrics_callbacks: Optional list of callbacks for adding additional
metrics to the graph (see EvalSharedModel for more information on how to
configure additional metrics). Metrics for example count and example
weights will be added automatically.
include_default_metrics: True to include the default metrics that are part
of the saved model graph during evaluation.
example_weight_key: Example weight key (single-output model) or dict of
example weight keys (multi-output model) keyed by output name.
additional_fetches: Prefixes of additional tensors stored in
signature_def.inputs that should be fetched at prediction time. The
"features" and "labels" tensors are handled automatically and should not
be included.
blacklist_feature_fetches: List of tensor names in the features dictionary
which should be excluded from the fetches request. This is useful in
scenarios where features are large (e.g. images) and can lead to excessive
memory use if stored.
"""
# Always compute example weight and example count.
# PyType doesn't know about the magic exports we do in post_export_metrics.
# Additionally, the lines seem to get reordered in compilation, so we can't
# just put the disable-attr on the add_metrics_callbacks lines.
# pytype: disable=module-attr
if not add_metrics_callbacks:
add_metrics_callbacks = []
example_count_callback = post_export_metrics.example_count()
add_metrics_callbacks.append(example_count_callback)
if example_weight_key:
if isinstance(example_weight_key, dict):
for output_name, key in example_weight_key.items():
example_weight_callback = post_export_metrics.example_weight(
key, metric_tag=output_name)
add_metrics_callbacks.append(example_weight_callback)
else:
example_weight_callback = post_export_metrics.example_weight(
example_weight_key)
add_metrics_callbacks.append(example_weight_callback)
# pytype: enable=module-attr
return types.EvalSharedModel(
model_path=eval_saved_model_path,
add_metrics_callbacks=add_metrics_callbacks,
include_default_metrics=include_default_metrics,
example_weight_key=example_weight_key,
additional_fetches=additional_fetches,
construct_fn=dofn.make_construct_fn(
eval_saved_model_path,
add_metrics_callbacks,
include_default_metrics,
additional_fetches=additional_fetches,
blacklist_feature_fetches=blacklist_feature_fetches))
