def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text) -> pipeline.Pipeline:
"""Implements the cifar10 pipeline with TFX."""
examples = external_input(data_root)
input_split = example_gen_pb2.Input(splits=[
example_gen_pb2.Input.Split(name='train', pattern='train.tfrecord'),
example_gen_pb2.Input.Split(name='eval', pattern='test.tfrecord')
])
example_gen = ImportExampleGen(input=examples, input_config=input_split)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(module_file=module_file,
examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=1000),
eval_args=trainer_pb2.EvalArgs(num_steps=500))
# Uses TFMA to compute a evaluation statistics over features of a model.
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(
specs=[evaluator_pb2.SingleSlicingSpec()]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats,
transform, trainer, evaluator, model_validator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
)
def _set_up_test_pipeline(self):
"""Builds an Iris example pipeline with slight changes."""
pipeline_name = "iris"
iris_root = "iris_root"
serving_model_dir = os.path.join(iris_root, "serving_model", pipeline_name)
tfx_root = "tfx_root"
data_path = os.path.join(tfx_root, "data_path")
pipeline_root = os.path.join(tfx_root, "pipelines", pipeline_name)
self.test_pipeline_info = data_types.PipelineInfo(pipeline_name, iris_root)
example_gen = CsvExampleGen(input=external_input(data_path))
statistics_gen = StatisticsGen(examples=example_gen.outputs["examples"])
schema_gen = SchemaGen(
statistics=statistics_gen.outputs["statistics"],
infer_feature_shape=True)
example_validator = ExampleValidator(
statistics=statistics_gen.outputs["statistics"],
schema=schema_gen.outputs["schema"])
trainer = Trainer(
# Use RuntimeParameter as module_file to test out RuntimeParameter in
# compiler.
module_file=data_types.RuntimeParameter(
name="module_file",
default=os.path.join(iris_root, "iris_utils.py"),
ptype=str),
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
examples=example_gen.outputs["examples"],
schema=schema_gen.outputs["schema"],
train_args=trainer_pb2.TrainArgs(num_steps=2000),
eval_args=trainer_pb2.EvalArgs(num_steps=5))
model_resolver = ResolverNode(
instance_name="latest_blessed_model_resolver",
resolver_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name="eval")],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
"sparse_categorical_accuracy":
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={"value": 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={"value": -1e-10}))
})
])
evaluator = Evaluator(
examples=example_gen.outputs["examples"],
model=trainer.outputs["model"],
baseline_model=model_resolver.outputs["model"],
eval_config=eval_config)
pusher = Pusher(
model=trainer.outputs["model"],
model_blessing=evaluator.outputs["blessing"],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
self._pipeline = pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen,
statistics_gen,
schema_gen,
example_validator,
trainer,
model_resolver,
evaluator,
pusher,
],
enable_cache=True,
beam_pipeline_args=[])
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input_base=data_root)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Get the latest model so that we can warm start from the model.
latest_model_resolver = ResolverNode(
instance_name='latest_model_resolver',
resolver_class=latest_artifacts_resolver.LatestArtifactsResolver,
latest_model=Channel(type=Model))
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
custom_executor_spec=executor_spec.ExecutorClassSpec(Executor),
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
base_model=latest_model_resolver.outputs['latest_model'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
# Change threshold will be ignored if there is no
# baseline model resolved from MLMD (first run).
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, schema_gen, example_validator,
transform, latest_model_resolver, trainer, model_resolver,
evaluator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=beam_pipeline_args)
def _create_parameterized_pipeline(
pipeline_name: Text, pipeline_root: Text, enable_cache: bool,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Creates a simple TFX pipeline with RuntimeParameter.
Args:
pipeline_name: The name of the pipeline.
pipeline_root: The root of the pipeline output.
enable_cache: Whether to enable cache in this pipeline.
beam_pipeline_args: Pipeline args for Beam jobs within Components.
Returns:
A logical TFX pipeline.Pipeline object.
"""
# First, define the pipeline parameters.
# Path to the CSV data file, under which there should be a data.csv file.
data_root = data_types.RuntimeParameter(
name='data-root',
default='gs://my-bucket/data',
ptype=Text,
)
# Path to the transform module file.
transform_module_file = data_types.RuntimeParameter(
name='transform-module',
default='gs://my-bucket/modules/transform_module.py',
ptype=Text,
)
# Path to the trainer module file.
trainer_module_file = data_types.RuntimeParameter(
name='trainer-module',
default='gs://my-bucket/modules/trainer_module.py',
ptype=Text,
)
# Number of epochs in training.
train_steps = data_types.RuntimeParameter(
name='train-steps',
default=10,
ptype=int,
)
# Number of epochs in evaluation.
eval_steps = data_types.RuntimeParameter(
name='eval-steps',
default=5,
ptype=int,
)
# The input data location is parameterized by data_root
examples = external_input(data_root)
example_gen = CsvExampleGen(input=examples)
statistics_gen = StatisticsGen(input_data=example_gen.outputs['examples'])
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# The module file used in Transform and Trainer component is paramterized by
# transform_module_file.
transform = Transform(input_data=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=transform_module_file)
# The numbers of steps in train_args are specified as RuntimeParameter with
# name 'train-steps' and 'eval-steps', respectively.
trainer = Trainer(
module_file=trainer_module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_output=transform.outputs['transform_graph'],
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps})
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
pusher = Pusher(
model_export=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=os.path.join(str(pipeline.ROOT_PARAMETER),
'model_serving'))))
return pipeline.Pipeline(pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, schema_gen,
example_validator, transform, trainer,
model_resolver, evaluator, pusher
],
enable_cache=enable_cache,
beam_pipeline_args=beam_pipeline_args)
def create_test_pipeline():
"""Builds an Iris example pipeline with slight changes."""
pipeline_name = "iris"
iris_root = "iris_root"
serving_model_dir = os.path.join(iris_root, "serving_model", pipeline_name)
tfx_root = "tfx_root"
data_path = os.path.join(tfx_root, "data_path")
pipeline_root = os.path.join(tfx_root, "pipelines", pipeline_name)
example_gen = CsvExampleGen(input_base=data_path)
statistics_gen = StatisticsGen(examples=example_gen.outputs["examples"])
my_importer = importer.Importer(
source_uri="m/y/u/r/i",
properties={
"split_names": "['train', 'eval']",
},
custom_properties={
"int_custom_property": 42,
"str_custom_property": "42",
},
artifact_type=standard_artifacts.Examples).with_id("my_importer")
schema_gen = SchemaGen(statistics=statistics_gen.outputs["statistics"],
infer_feature_shape=True)
example_validator = ExampleValidator(
statistics=statistics_gen.outputs["statistics"],
schema=schema_gen.outputs["schema"])
trainer = Trainer(
# Use RuntimeParameter as module_file to test out RuntimeParameter in
# compiler.
module_file=data_types.RuntimeParameter(name="module_file",
default=os.path.join(
iris_root,
"iris_utils.py"),
ptype=str),
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
examples=example_gen.outputs["examples"],
schema=schema_gen.outputs["schema"],
train_args=trainer_pb2.TrainArgs(num_steps=2000),
# Attaching `TrainerArgs` as platform config is not sensible practice,
# but is only for testing purpose.
eval_args=trainer_pb2.EvalArgs(num_steps=5)).with_platform_config(
config=trainer_pb2.TrainArgs(num_steps=2000))
model_resolver = resolver.Resolver(
strategy_class=latest_blessed_model_strategy.
LatestBlessedModelStrategy,
baseline_model=Channel(type=standard_artifacts.Model,
producer_component_id="Trainer"),
# Cannot add producer_component_id="Evaluator" for model_blessing as it
# raises "producer component should have already been compiled" error.
model_blessing=Channel(type=standard_artifacts.ModelBlessing)).with_id(
"latest_blessed_model_resolver")
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name="eval")],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
"sparse_categorical_accuracy":
tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={"value": 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={"value": -1e-10}))
})
])
evaluator = Evaluator(
examples=example_gen.outputs["examples"],
model=trainer.outputs["model"],
baseline_model=model_resolver.outputs["baseline_model"],
eval_config=eval_config)
pusher = Pusher(model=trainer.outputs["model"],
model_blessing=evaluator.outputs["blessing"],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen,
statistics_gen,
my_importer,
schema_gen,
example_validator,
trainer,
model_resolver,
evaluator,
pusher,
],
enable_cache=False,
beam_pipeline_args=["--my_testing_beam_pipeline_args=bar"],
# Attaching `TrainerArgs` as platform config is not sensible practice,
# but is only for testing purpose.
platform_config=trainer_pb2.TrainArgs(num_steps=2000),
execution_mode=pipeline.ExecutionMode.ASYNC)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text, enable_tuning: bool,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Implements the Iris flowers pipeline with TFX."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Tunes the hyperparameters for model training based on user-provided Python
# function. Note that once the hyperparameters are tuned, you can drop the
# Tuner component from pipeline and feed Trainer with tuned hyperparameters.
if enable_tuning:
tuner = Tuner(module_file=module_file,
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=20),
eval_args=trainer_pb2.EvalArgs(num_steps=5))
# Uses user-provided Python function that trains a model.
trainer = Trainer(
module_file=module_file,
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
# If Tuner is in the pipeline, Trainer can take Tuner's output
# best_hyperparameters artifact as input and utilize it in the user module
# code.
#
# If there isn't Tuner in the pipeline, either use ImporterNode to import
# a previous Tuner's output to feed to Trainer, or directly use the tuned
# hyperparameters in user module code and set hyperparameters to None
# here.
#
# Example of ImporterNode,
# hparams_importer = ImporterNode(
# instance_name='import_hparams',
# source_uri='path/to/best_hyperparameters.txt',
# artifact_type=HyperParameters)
# ...
# hyperparameters = hparams_importer.outputs['result'],
hyperparameters=(tuner.outputs['best_hyperparameters']
if enable_tuning else None),
train_args=trainer_pb2.TrainArgs(num_steps=100),
eval_args=trainer_pb2.EvalArgs(num_steps=5))
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute an evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='variety')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='SparseCategoricalAccuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10})))
])
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
components = [
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
trainer,
model_resolver,
evaluator,
pusher,
]
if enable_tuning:
components.append(tuner)
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=beam_pipeline_args)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
trainer_module_file: Text, evaluator_module_file: Text,
serving_model_dir: Text, metadata_path: Text,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Implements the Penguin pipeline with TFX."""
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input_base=data_root)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# TODO(humichael): Handle applying transformation component in Milestone 3.
# Uses user-provided Python function that trains a model using TF-Learn.
# Num_steps is not provided during evaluation because the scikit-learn model
# loads and evaluates the entire test set at once.
trainer = Trainer(module_file=trainer_module_file,
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
train_args=trainer_pb2.TrainArgs(num_steps=2000),
eval_args=trainer_pb2.EvalArgs())
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(
type=ModelBlessing)).with_id('latest_blessed_model_resolver')
# Uses TFMA to compute evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='species')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='Accuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10})))
])
])
evaluator = Evaluator(module_file=evaluator_module_file,
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen,
statistics_gen,
schema_gen,
example_validator,
trainer,
model_resolver,
evaluator,
pusher,
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=beam_pipeline_args,
)
def create_tfx_pipeline(pipeline_name: Text, input_dir: Text):
# Output 2 splits: train:eval=3:1.
example_gen = CsvExampleGen(input=external_input(RAW_DATA))
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Fetch the latest trained model under the same context for warm-starting.
latest_model_resolver = ResolverNode(
instance_name='latest_model_resolver',
resolver_class=latest_artifacts_resolver.LatestArtifactsResolver,
model=channel.Channel(type=standard_artifacts.Model))
trainer = Trainer(
module_file=TRAINER_FILE,
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
train_args=trainer_pb2.TrainArgs(num_steps=3000),
eval_args=trainer_pb2.EvalArgs(num_steps=3000))
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=channel.Channel(type=standard_artifacts.Model),
model_blessing=channel.Channel(type=standard_artifacts.ModelBlessing))
# Set the TFMA config for Model Evaluation and Validation.
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='label')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='SparseCategoricalAccuracy',
threshold=tfma.MetricThreshold(
# Accept models only if SparseCategoricalAccuracy > 0.8
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.8}),
# TODO: modify this
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-2})))
])
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=SERVING_MODEL_DIR)))
components = [
example_gen, statistics_gen, schema_gen, example_validator, transform,
trainer, pusher
]
return tfx_pipeline.Pipeline(pipeline_name=pipeline_name,
pipeline_root=PIPELINE_ROOT,
components=components)
def _create_pipeline(
pipeline_name: Text, pipeline_root: Text, query: Text, module_file: Text,
beam_pipeline_args: List[Text], ai_platform_training_args: Dict[Text, Text],
bigquery_serving_args: Dict[Text, Text]) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX and Kubeflow Pipelines."""
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = big_query_example_gen_component.BigQueryExampleGen(query=query)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(
statistics=statistics_gen.outputs['statistics'], infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn
# to train a model on Google Cloud AI Platform.
trainer = Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.Executor),
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000),
custom_config={'ai_platform_training_args': ai_platform_training_args})
# Uses TFMA to compute a evaluation statistics over features of a model.
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(
examples=example_gen.outputs['examples'], model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to Google Cloud BigQuery ML if check passed.
pusher = Pusher(
custom_executor_spec=executor_spec.ExecutorClassSpec(
bigquery_ml_pusher_executor.Executor),
model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
custom_config={'bigquery_serving_args': bigquery_serving_args})
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, schema_gen, example_validator, transform,
trainer, evaluator, model_validator, pusher
],
beam_pipeline_args=beam_pipeline_args,
)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_path: Text,
# TODO(step 7): (Optional) Uncomment here to use BigQuery as a data source.
# query: Text,
preprocessing_fn: Text,
trainer_fn: Text,
train_args: trainer_pb2.TrainArgs,
eval_args: trainer_pb2.EvalArgs,
serving_model_dir: Text,
metadata_connection_config: Optional[
metadata_store_pb2.ConnectionConfig] = None,
beam_pipeline_args: Optional[List[Text]] = None,
ai_platform_training_args: Optional[Dict[Text, Text]] = None,
ai_platform_serving_args: Optional[Dict[Text, Any]] = None,
) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
components = []
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=external_input(data_path))
# TODO(step 7): (Optional) Uncomment here to use BigQuery as a data source.
# example_gen = BigQueryExampleGen(query=query)
components.append(example_gen)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# TODO(step 5): Uncomment here to add StatisticsGen to the pipeline.
components.append(statistics_gen)
# Generates schema based on statistics files.
infer_schema = SchemaGen(
statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# TODO(step 5): Uncomment here to add SchemaGen to the pipeline.
components.append(infer_schema)
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator( # pylint: disable=unused-variable
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# TODO(step 5): Uncomment here to add ExampleValidator to the pipeline.
components.append(validate_stats)
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
preprocessing_fn=preprocessing_fn)
# TODO(step 6): Uncomment here to add Transform to the pipeline.
components.append(transform)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer_args = {
'trainer_fn': trainer_fn,
'transformed_examples': transform.outputs['transformed_examples'],
'schema': infer_schema.outputs['schema'],
'transform_graph': transform.outputs['transform_graph'],
'train_args': train_args,
'eval_args': eval_args,
}
if ai_platform_training_args is not None:
trainer_args.update({
'custom_executor_spec':
executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.Executor),
'custom_config': {
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
ai_platform_training_args
}
})
trainer = Trainer(**trainer_args)
# TODO(step 6): Uncomment here to add Trainer to the pipeline.
components.append(trainer)
# Uses TFMA to compute a evaluation statistics over features of a model.
model_analyzer = Evaluator( # pylint: disable=unused-variable
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# TODO(step 6): Uncomment here to add Evaluator to the pipeline.
components.append(model_analyzer)
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(
examples=example_gen.outputs['examples'], model=trainer.outputs['model'])
# TODO(step 6): Uncomment here to add ModelValidator to the pipeline.
components.append(model_validator)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher_args = {
'model':
trainer.outputs['model'],
'model_blessing':
model_validator.outputs['blessing'],
'push_destination':
pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)),
}
if ai_platform_serving_args is not None:
pusher_args.update({
'custom_executor_spec':
executor_spec.ExecutorClassSpec(ai_platform_pusher_executor.Executor
),
'custom_config': {
ai_platform_pusher_executor.SERVING_ARGS_KEY:
ai_platform_serving_args
},
})
pusher = Pusher(**pusher_args) # pylint: disable=unused-variable
# TODO(step 6): Uncomment here to add Pusher to the pipeline.
components.append(pusher)
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
enable_cache=False,
metadata_connection_config=metadata_connection_config,
beam_pipeline_args=beam_pipeline_args,
)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text,
worker_parallelism: int,
runner: Text) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(
statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Uses TFMA to compute a evaluation statistics over features of a model.
model_analyzer = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(
examples=example_gen.outputs['examples'], model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(
model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
beam_common_args = [
'--environment_type=LOOPBACK',
'--sdk_worker_parallelism=%d' % worker_parallelism,
'--experiments=use_loopback_process_worker=True',
# Setting environment_cache_millis to practically infinity enables
# continual reuse of Beam SDK workers, improving performance.
'--environment_cache_millis=1000000',
# Note; We use 100 worker threads to mitigate the issue with
# scheduling work between the Beam runner and SDK harness. Flink
# and Spark can process unlimited work items concurrently while
# SdkHarness can only process 1 work item per worker thread.
# Having 100 threads will let 100 tasks execute concurrently
# avoiding scheduling issue in most cases. In case the threads are
# exhausted, beam prints the relevant message in the log.
# TODO(BEAM-8151) Remove worker_threads=100 after we start using a # pylint: disable=g-bad-todo
# virtually unlimited thread pool by default.
'--experiments=worker_threads=100',
]
if runner == 'flink':
beam_pipeline_args = beam_common_args + [
'--runner=FlinkRunner',
'--flink_master=localhost:8081',
'--flink_submit_uber_jar',
# TODO(ibzib) move these to flink.conf
'--parallelism=%d' % worker_parallelism,
# TODO(FLINK-10672): Obviate setting BATCH_FORCED. # pylint: disable=g-bad-todo
'--execution_mode_for_batch=BATCH_FORCED',
]
elif runner == 'spark':
beam_pipeline_args = beam_common_args + [
# TODO(ibzib)
]
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats, transform,
trainer, model_analyzer, model_validator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
# LINT.IfChange
beam_pipeline_args=beam_pipeline_args,
# LINT.ThenChange(setup/setup_beam_on_spark.sh)
# LINT.ThenChange(setup/setup_beam_on_flink.sh)
)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_path: Text,
preprocessing_fn: Text,
run_fn: Text,
train_args: trainer_pb2.TrainArgs,
eval_args: trainer_pb2.EvalArgs,
eval_accuracy_threshold: float,
serving_model_dir: Text,
metadata_connection_config: Optional[
metadata_store_pb2.ConnectionConfig
] = None,
beam_pipeline_args: Optional[List[Text]] = None,
ai_platform_training_args: Optional[Dict[Text, Text]] = None,
ai_platform_serving_args: Optional[Dict[Text, Any]] = None,
) -> pipeline.Pipeline:
"""Implements the complaint prediction pipeline with TFX."""
components = []
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input_base=data_path)
components.append(example_gen)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs["examples"])
components.append(statistics_gen)
# Generates schema based on statistics files.
schema_gen = SchemaGen(
statistics=statistics_gen.outputs["statistics"],
infer_feature_shape=True,
)
components.append(schema_gen)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator( # pylint: disable=unused-variable
statistics=statistics_gen.outputs["statistics"],
schema=schema_gen.outputs["schema"],
)
components.append(example_validator)
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs["examples"],
schema=schema_gen.outputs["schema"],
preprocessing_fn=preprocessing_fn,
)
components.append(transform)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer_args = {
"run_fn": run_fn,
"transformed_examples": transform.outputs["transformed_examples"],
"schema": schema_gen.outputs["schema"],
"transform_graph": transform.outputs["transform_graph"],
"train_args": train_args,
"eval_args": eval_args,
"custom_executor_spec": executor_spec.ExecutorClassSpec(
trainer_executor.GenericExecutor
),
}
if ai_platform_training_args is not None:
trainer_args.update(
{
"custom_executor_spec": executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.GenericExecutor
),
"custom_config": {
ai_platform_trainer_executor.TRAINING_ARGS_KEY: ai_platform_training_args, # noqa
},
}
)
trainer = Trainer(**trainer_args)
components.append(trainer)
# Get the latest blessed model for model validation.
model_resolver = tfx.dsl.Resolver(
strategy_class=tfx.dsl.experimental.LatestBlessedModelStrategy,
model=tfx.dsl.Channel(type=tfx.types.standard_artifacts.Model),
model_blessing=tfx.dsl.Channel(
type=tfx.types.standard_artifacts.ModelBlessing)).with_id(
'latest_blessed_model_resolver')
components.append(model_resolver)
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key="big_tipper")],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(
metrics=[
tfma.MetricConfig(
class_name="BinaryAccuracy",
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={"value": eval_accuracy_threshold}
),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={"value": -1e-10},
),
),
)
]
)
],
)
evaluator = Evaluator(
examples=example_gen.outputs["examples"],
model=trainer.outputs["model"],
baseline_model=model_resolver.outputs["model"],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config,
)
components.append(evaluator)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher_args = {
"model": trainer.outputs["model"],
"model_blessing": evaluator.outputs["blessing"],
"push_destination": pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir
)
),
}
if ai_platform_serving_args is not None:
pusher_args.update(
{
"custom_executor_spec": executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor
),
"custom_config": {
ai_platform_pusher_executor.SERVING_ARGS_KEY: ai_platform_serving_args # noqa
},
}
)
pusher = Pusher(**pusher_args) # pylint: disable=unused-variable
components.append(pusher)
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
enable_cache=True,
metadata_connection_config=metadata_connection_config,
beam_pipeline_args=beam_pipeline_args,
)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir_lite: Text,
metadata_path: Text, labels_path: Text,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Implements the CIFAR10 image classification pipeline using TFX."""
# This is needed for datasets with pre-defined splits
# Change the pattern argument to train_whole/* and test_whole/* to train
# on the whole CIFAR-10 dataset
input_config = example_gen_pb2.Input(splits=[
example_gen_pb2.Input.Split(name='train', pattern='train/*'),
example_gen_pb2.Input.Split(name='eval', pattern='test/*')
])
# Brings data into the pipeline.
example_gen = ImportExampleGen(input_base=data_root,
input_config=input_config)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that trains a model.
# When traning on the whole dataset, use 18744 for train steps, 156 for eval
# steps. 18744 train steps correspond to 24 epochs on the whole train set, and
# 156 eval steps correspond to 1 epoch on the whole test set. The
# configuration below is for training on the dataset we provided in the data
# folder, which has 128 train and 128 test samples. The 160 train steps
# correspond to 40 epochs on this tiny train set, and 4 eval steps correspond
# to 1 epoch on this tiny test set.
trainer = Trainer(module_file=module_file,
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
train_args=trainer_pb2.TrainArgs(num_steps=160),
eval_args=trainer_pb2.EvalArgs(num_steps=4),
custom_config={'labels_path': labels_path})
# Get the latest blessed model for model validation.
model_resolver = resolver.Resolver(
strategy_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(
type=ModelBlessing)).with_id('latest_blessed_model_resolver')
# Uses TFMA to compute evaluation statistics over features of a model and
# perform quality validation of a candidate model (compare to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[
tfma.ModelSpec(label_key='label_xf', model_type='tf_lite')
],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='SparseCategoricalAccuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.55}),
# Change threshold will be ignored if there is no
# baseline model resolved from MLMD (first run).
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-3})))
])
])
# Uses TFMA to compute the evaluation statistics over features of a model.
# We evaluate using the materialized examples that are output by Transform
# because
# 1. the decoding_png function currently performed within Transform are not
# compatible with TFLite.
# 2. MLKit requires deserialized (float32) tensor image inputs
# Note that for deployment, the same logic that is performed within Transform
# must be reproduced client-side.
evaluator = Evaluator(examples=transform.outputs['transformed_examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir_lite)))
components = [
example_gen, statistics_gen, schema_gen, example_validator, transform,
trainer, model_resolver, evaluator, pusher
]
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=beam_pipeline_args)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Uses TFMA to compute a evaluation statistics over features of a model.
model_analyzer = Evaluator(
examples=example_gen.outputs['examples'],
model_exports=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats,
transform, trainer, model_analyzer, model_validator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
# LINT.IfChange
beam_pipeline_args=[
# ----- Beam Args -----.
'--runner=PortableRunner',
# Points to the job server started in
# setup_beam_on_(flink|spark).sh
'--job_endpoint=localhost:8099',
'--environment_type=LOOPBACK',
# TODO(BEAM-6754): Utilize multicore in LOOPBACK environment. # pylint: disable=g-bad-todo
# TODO(BEAM-5167): Use concurrency information from SDK Harness. # pylint: disable=g-bad-todo
# Note; We use 100 worker threads to mitigate the issue with
# scheduling work between the Beam runner and SDK harness. Flink
# and Spark can process unlimited work items concurrently while
# SdkHarness can only process 1 work item per worker thread.
# Having 100 threads will let 100 tasks execute concurrently
# avoiding scheduling issue in most cases. In case the threads are
# exhausted, beam print the relevant message in the log.
'--experiments=worker_threads=100',
# TODO(BEAM-7199): Obviate the need for setting pre_optimize=all. # pylint: disable=g-bad-todo
'--experiments=pre_optimize=all',
# ----- Flink runner-specific Args -----.
# TODO(b/126725506): Set the task parallelism based on cpu cores.
# TODO(FLINK-10672): Obviate setting BATCH_FORCED.
'--execution_mode_for_batch=BATCH_FORCED',
],
# LINT.ThenChange(setup/setup_beam_on_spark.sh)
# LINT.ThenChange(setup/setup_beam_on_flink.sh)
)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_path: Text,
# TODO(step 7): (Optional) Uncomment here to use BigQuery as a data source.
# query: Text,
preprocessing_fn: Text,
run_fn: Text,
train_args: trainer_pb2.TrainArgs,
eval_args: trainer_pb2.EvalArgs,
eval_accuracy_threshold: float,
serving_model_dir: Text,
metadata_connection_config: Optional[
metadata_store_pb2.ConnectionConfig] = None,
beam_pipeline_args: Optional[List[Text]] = None,
ai_platform_training_args: Optional[Dict[Text, Text]] = None,
ai_platform_serving_args: Optional[Dict[Text, Any]] = None,
) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
components = []
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=external_input(data_path))
# TODO(step 7): (Optional) Uncomment here to use BigQuery as a data source.
# example_gen = big_query_example_gen_component.BigQueryExampleGen(
# query=query)
components.append(example_gen)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# TODO(step 5): Uncomment here to add StatisticsGen to the pipeline.
# components.append(statistics_gen)
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# TODO(step 5): Uncomment here to add SchemaGen to the pipeline.
# components.append(schema_gen)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator( # pylint: disable=unused-variable
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# TODO(step 5): Uncomment here to add ExampleValidator to the pipeline.
# components.append(example_validator)
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
preprocessing_fn=preprocessing_fn)
# TODO(step 6): Uncomment here to add Transform to the pipeline.
# components.append(transform)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer_args = {
'run_fn': run_fn,
'transformed_examples': transform.outputs['transformed_examples'],
'schema': schema_gen.outputs['schema'],
'transform_graph': transform.outputs['transform_graph'],
'train_args': train_args,
'eval_args': eval_args,
}
if ai_platform_training_args is not None:
trainer_args.update({
'custom_executor_spec':
executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.GenericExecutor),
'custom_config': {
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
ai_platform_training_args,
}
})
trainer = Trainer(**trainer_args)
# TODO(step 6): Uncomment here to add Trainer to the pipeline.
# components.append(trainer)
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# TODO(step 6): Uncomment here to add ResolverNode to the pipeline.
# components.append(model_resolver)
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='big_tipper')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='BinaryAccuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': eval_accuracy_threshold}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10})))
])
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# TODO(step 6): Uncomment here to add Evaluator to the pipeline.
# components.append(evaluator)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher_args = {
'model':
trainer.outputs['model'],
'model_blessing':
evaluator.outputs['blessing'],
'push_destination':
pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)),
}
if ai_platform_serving_args is not None:
pusher_args.update({
'custom_executor_spec':
executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
'custom_config': {
ai_platform_pusher_executor.SERVING_ARGS_KEY:
ai_platform_serving_args
},
})
pusher = Pusher(**pusher_args) # pylint: disable=unused-variable
# TODO(step 6): Uncomment here to add Pusher to the pipeline.
# components.append(pusher)
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
# Change this value to control caching of execution results. Default value
# is `False`.
# enable_cache=True,
metadata_connection_config=metadata_connection_config,
beam_pipeline_args=beam_pipeline_args,
)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text,
direct_num_workers: int) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=module_file)
# Get the latest model so that we can warm start from the model.
latest_model_resolver = ResolverNode(
instance_name='latest_model_resolver',
resolver_class=latest_artifacts_resolver.LatestArtifactsResolver,
latest_model=Channel(type=Model))
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
base_model=latest_model_resolver.outputs['latest_model'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Uses TFMA to compute a evaluation statistics over features of a model.
model_analyzer = Evaluator(
examples=example_gen.outputs['examples'],
model_exports=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats,
transform, latest_model_resolver, trainer, model_analyzer,
model_validator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
# TODO(b/142684737): The multi-processing API might change.
beam_pipeline_args=['--direct_num_workers=%d' % direct_num_workers])
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
direct_num_workers: int) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX and Kubeflow Pipelines."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn
# to train a model on Google Cloud AI Platform.
trainer = Trainer(
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000),
)
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# Performs infra validation of a candidate model to prevent unservable model
# from being pushed. In order to use InfraValidator component, persistent
# volume and its claim that the pipeline is using should be a ReadWriteMany
# access mode.
infra_validator = InfraValidator(
model=trainer.outputs['model'],
examples=example_gen.outputs['examples'],
serving_spec=infra_validator_pb2.ServingSpec(
tensorflow_serving=infra_validator_pb2.TensorFlowServing(
tags=['latest']),
kubernetes=infra_validator_pb2.KubernetesConfig()),
request_spec=infra_validator_pb2.RequestSpec(
tensorflow_serving=infra_validator_pb2.
TensorFlowServingRequestSpec()))
# Checks whether the model passed the validation steps and pushes the model
# to Google Cloud AI Platform if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
infra_blessing=infra_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
trainer,
model_resolver,
evaluator,
infra_validator,
pusher,
],
# TODO(b/142684737): The multi-processing API might change.
beam_pipeline_args=['--direct_num_workers=%d' % direct_num_workers],
)
def testTaxiPipelineNewStyleCompatibility(self):
example_gen = CsvExampleGen(input_base='/tmp/fake/path')
statistics_gen = StatisticsGen(
examples=example_gen.outputs['examples'])
self.assertIs(statistics_gen.inputs['examples'],
statistics_gen.inputs['input_data'])
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'])
self.assertIs(schema_gen.inputs['statistics'],
schema_gen.inputs['stats'])
self.assertIs(schema_gen.outputs['schema'],
schema_gen.outputs['output'])
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file='/tmp/fake/module/file')
self.assertIs(transform.inputs['examples'],
transform.inputs['input_data'])
self.assertIs(transform.outputs['transform_graph'],
transform.outputs['transform_output'])
trainer = Trainer(
module_file='/tmp/fake/module/file',
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
self.assertIs(trainer.inputs['transform_graph'],
trainer.inputs['transform_output'])
self.assertIs(trainer.outputs['model'], trainer.outputs['output'])
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.
HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
self.assertIs(evaluator.inputs['model'],
evaluator.inputs['model_exports'])
self.assertIs(evaluator.outputs['evaluation'],
evaluator.outputs['output'])
pusher = Pusher(model=trainer.outputs['output'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory='/fake/serving/dir')))
self.assertIs(pusher.inputs['model'], pusher.inputs['model_export'])
self.assertIs(pusher.outputs['pushed_model'],
pusher.outputs['model_push'])
def create_pipeline(pipeline_name: Text,
pipeline_root: Text,
data_root_uri: data_types.RuntimeParameter,
train_steps: data_types.RuntimeParameter,
eval_steps: data_types.RuntimeParameter,
enable_tuning: bool,
ai_platform_training_args: Dict[Text, Text],
ai_platform_serving_args: Dict[Text, Text],
beam_pipeline_args: List[Text],
enable_cache: Optional[bool] = False) -> pipeline.Pipeline:
"""Trains and deploys the Keras Covertype Classifier with TFX and Kubeflow Pipeline on Google Cloud.
Args:
pipeline_name: name of the TFX pipeline being created.
pipeline_root: root directory of the pipeline. Should be a valid GCS path.
data_root_uri: uri of the dataset.
train_steps: runtime parameter for number of model training steps for the Trainer component.
eval_steps: runtime parameter for number of model evaluation steps for the Trainer component.
enable_tuning: If True, the hyperparameter tuning through CloudTuner is
enabled.
ai_platform_training_args: Args of CAIP training job. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#Job
for detailed description.
ai_platform_serving_args: Args of CAIP model deployment. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models
for detailed description.
beam_pipeline_args: Optional list of beam pipeline options. Please refer to
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options.
When this argument is not provided, the default is to use GCP
DataflowRunner with 50GB disk size as specified in this function. If an
empty list is passed in, default specified by Beam will be used, which can
be found at
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options
enable_cache: Optional boolean
Returns:
A TFX pipeline object.
"""
# Brings data into the pipeline and splits the data into training and eval splits
output_config = example_gen_pb2.Output(
split_config=example_gen_pb2.SplitConfig(splits=[
example_gen_pb2.SplitConfig.Split(name='train', hash_buckets=4),
example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1)
]))
examplegen = CsvExampleGen(input_base=data_root_uri)
# Computes statistics over data for visualization and example validation.
statisticsgen = StatisticsGen(examples=examplegen.outputs.examples)
# Generates schema based on statistics files. Even though, we use user-provided schema
# we still want to generate the schema of the newest data for tracking and comparison
schemagen = SchemaGen(statistics=statisticsgen.outputs.statistics)
# Import a user-provided schema
import_schema = ImporterNode(instance_name='import_user_schema',
source_uri=SCHEMA_FOLDER,
artifact_type=Schema)
# Performs anomaly detection based on statistics and data schema.
examplevalidator = ExampleValidator(
statistics=statisticsgen.outputs.statistics,
schema=import_schema.outputs.result)
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=examplegen.outputs.examples,
schema=import_schema.outputs.result,
module_file=TRANSFORM_MODULE_FILE)
# Tunes the hyperparameters for model training based on user-provided Python
# function. Note that once the hyperparameters are tuned, you can drop the
# Tuner component from pipeline and feed Trainer with tuned hyperparameters.
if enable_tuning:
# The Tuner component launches 1 AI Platform Training job for flock management.
# For example, n_workers (defined by num_parallel_trials) in the flock
# management AI Platform Training job, each run Tuner.Executor in parallel.
tuner = Tuner(
module_file=TRAIN_MODULE_FILE,
examples=transform.outputs.transformed_examples,
transform_graph=transform.outputs.transform_graph,
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps},
tune_args=tuner_pb2.TuneArgs(
# num_parallel_trials can be configured for distributed training.
num_parallel_trials=1),
custom_config={
# Configures Cloud AI Platform-specific configs. For details, see
# https://cloud.google.com/ai-platform/training/docs/reference/rest/v1/projects.jobs#traininginput.
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
ai_platform_training_args
})
# Trains the model using a user provided trainer function.
trainer = Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.GenericExecutor),
module_file=TRAIN_MODULE_FILE,
transformed_examples=transform.outputs.transformed_examples,
schema=import_schema.outputs.result,
transform_graph=transform.outputs.transform_graph,
hyperparameters=(tuner.outputs.best_hyperparameters
if enable_tuning else None),
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps},
custom_config={'ai_platform_training_args': ai_platform_training_args})
# Get the latest blessed model for model validation.
resolver = ResolverNode(instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute a evaluation statistics over features of a model.
accuracy_threshold = tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(lower_bound={'value': 0.5},
upper_bound={'value':
0.99}), )
metrics_specs = tfma.MetricsSpec(metrics=[
tfma.MetricConfig(class_name='SparseCategoricalAccuracy',
threshold=accuracy_threshold),
tfma.MetricConfig(class_name='ExampleCount')
])
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='Cover_Type')],
metrics_specs=[metrics_specs],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['Wilderness_Area'])
])
evaluator = Evaluator(examples=examplegen.outputs.examples,
model=trainer.outputs.model,
baseline_model=resolver.outputs.model,
eval_config=eval_config)
# Validate model can be loaded and queried in sand-boxed environment
# mirroring production.
serving_config = infra_validator_pb2.ServingSpec(
tensorflow_serving=infra_validator_pb2.TensorFlowServing(
tags=['latest']),
kubernetes=infra_validator_pb2.KubernetesConfig(),
)
validation_config = infra_validator_pb2.ValidationSpec(
max_loading_time_seconds=60,
num_tries=3,
)
request_config = infra_validator_pb2.RequestSpec(
tensorflow_serving=infra_validator_pb2.TensorFlowServingRequestSpec(),
num_examples=3,
)
infravalidator = InfraValidator(
model=trainer.outputs.model,
examples=examplegen.outputs.examples,
serving_spec=serving_config,
validation_spec=validation_config,
request_spec=request_config,
)
# Checks whether the model passed the validation steps and pushes the model
# to CAIP Prediction if checks are passed.
pusher = Pusher(custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
model=trainer.outputs.model,
model_blessing=evaluator.outputs.blessing,
infra_blessing=infravalidator.outputs.blessing,
custom_config={
ai_platform_pusher_executor.SERVING_ARGS_KEY:
ai_platform_serving_args
})
components = [
examplegen, statisticsgen, schemagen, import_schema, examplevalidator,
transform, trainer, resolver, evaluator, infravalidator, pusher
]
if enable_tuning:
components.append(tuner)
return pipeline.Pipeline(pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
enable_cache=enable_cache,
beam_pipeline_args=beam_pipeline_args)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text,
training_data_root: Text, inference_data_root: Text,
module_file: Text, metadata_path: Text,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
# Brings training data into the pipeline or otherwise joins/converts
# training data.
training_example_gen = CsvExampleGen(
input_base=training_data_root).with_id('training_example_gen')
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(
examples=training_example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=training_example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
custom_executor_spec=executor_spec.ExecutorClassSpec(Executor),
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Get the latest blessed model for model validation.
model_resolver = resolver.Resolver(
strategy_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(
type=ModelBlessing)).with_id('latest_blessed_model_resolver')
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
# Change threshold will be ignored if there is no
# baseline model resolved from MLMD (first run).
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(examples=training_example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
# Brings inference data into the pipeline.
inference_example_gen = CsvExampleGen(
input_base=inference_data_root,
output_config=example_gen_pb2.Output(
split_config=example_gen_pb2.SplitConfig(splits=[
example_gen_pb2.SplitConfig.Split(name='unlabelled',
hash_buckets=100)
]))).with_id('inference_example_gen')
# Performs offline batch inference over inference examples.
bulk_inferrer = BulkInferrer(
examples=inference_example_gen.outputs['examples'],
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
# Empty data_spec.example_splits will result in using all splits.
data_spec=bulk_inferrer_pb2.DataSpec(),
model_spec=bulk_inferrer_pb2.ModelSpec())
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
training_example_gen, inference_example_gen, statistics_gen,
schema_gen, example_validator, transform, trainer, model_resolver,
evaluator, bulk_inferrer
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=beam_pipeline_args)
def _create_pipeline(pipeline_name: str, pipeline_root: str, data_root: str,
module_file: str, serving_model_dir: str,
metadata_path: str,
beam_pipeline_args: List[str]) -> pipeline.Pipeline:
"""Implements the Bert classication on Cola dataset pipline with TFX."""
input_config = example_gen_pb2.Input(splits=[
example_gen_pb2.Input.Split(name='train', pattern='train/*'),
example_gen_pb2.Input.Split(name='eval', pattern='validation/*')
])
# Brings data into the pipline
example_gen = CsvExampleGen(input_base=data_root,
input_config=input_config)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that trains a model.
trainer = Trainer(
module_file=module_file,
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
# Adjust these steps when training on the full dataset.
train_args=trainer_pb2.TrainArgs(num_steps=2),
eval_args=trainer_pb2.EvalArgs(num_steps=1))
# Get the latest blessed model for model validation.
model_resolver = resolver.Resolver(
strategy_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(
type=ModelBlessing)).with_id('latest_blessed_model_resolver')
# Uses TFMA to compute evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='label')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='SparseCategoricalAccuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
# Adjust the threshold when training on the
# full dataset.
lower_bound={'value': 0.5}),
# Change threshold will be ignored if there is no
# baseline model resolved from MLMD (first run).
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-2})))
])
])
evaluator = Evaluator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
components = [
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
trainer,
model_resolver,
evaluator,
pusher,
]
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
enable_cache=True,
beam_pipeline_args=beam_pipeline_args,
)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_root: Text,
module_file: Text,
ai_platform_training_args: Dict[Text, Text],
pusher_custom_config: Dict[Text, Text],
enable_tuning: bool,
beam_pipeline_args: Optional[List[Text]] = None) -> pipeline.Pipeline:
"""Implements the penguin pipeline with TFX and Kubeflow Pipeline.
Args:
pipeline_name: name of the TFX pipeline being created.
pipeline_root: root directory of the pipeline. Should be a valid GCS path.
data_root: uri of the penguin data.
module_file: uri of the module files used in Trainer and Transform
components.
ai_platform_training_args: Args of CAIP training job. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#Job
for detailed description.
pusher_custom_config: Custom configs passed to pusher.
enable_tuning: If True, the hyperparameter tuning through CloudTuner is
enabled.
beam_pipeline_args: Optional list of beam pipeline options. Please refer to
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options.
When this argument is not provided, the default is to use GCP
DataflowRunner with 50GB disk size as specified in this function. If an
empty list is passed in, default specified by Beam will be used, which can
be found at
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options
Returns:
A TFX pipeline object.
"""
# Beam args to run data processing on DataflowRunner.
#
# TODO(b/151114974): Remove `disk_size_gb` flag after default is increased.
# TODO(b/151116587): Remove `shuffle_mode` flag after default is changed.
# TODO(b/156874687): Remove `machine_type` after IP addresses are no longer a
# scaling bottleneck.
if beam_pipeline_args is None:
beam_pipeline_args = [
'--runner=DataflowRunner',
'--project=' + _project_id,
'--temp_location=' + os.path.join(_output_bucket, 'tmp'),
'--region=' + _gcp_region,
# Temporary overrides of defaults.
'--disk_size_gb=50',
'--experiments=shuffle_mode=auto',
'--machine_type=e2-standard-8',
]
# Number of epochs in training.
train_steps = data_types.RuntimeParameter(
name='train_steps',
default=100,
ptype=int,
)
# Number of epochs in evaluation.
eval_steps = data_types.RuntimeParameter(
name='eval_steps',
default=50,
ptype=int,
)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input_base=data_root)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(
statistics=statistics_gen.outputs['statistics'], infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Update ai_platform_training_args if distributed training was enabled.
# Number of worker machines used in distributed training.
worker_count = data_types.RuntimeParameter(
name='worker_count',
default=2,
ptype=int,
)
# Type of worker machines used in distributed training.
worker_type = data_types.RuntimeParameter(
name='worker_type',
default='standard',
ptype=str,
)
local_training_args = copy.deepcopy(ai_platform_training_args)
if FLAGS.distributed_training:
local_training_args.update({
# You can specify the machine types, the number of replicas for workers
# and parameter servers.
# https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#ScaleTier
'scaleTier': 'CUSTOM',
'masterType': 'large_model',
'workerType': worker_type,
'parameterServerType': 'standard',
'workerCount': worker_count,
'parameterServerCount': 1,
})
# Tunes the hyperparameters for model training based on user-provided Python
# function. Note that once the hyperparameters are tuned, you can drop the
# Tuner component from pipeline and feed Trainer with tuned hyperparameters.
if enable_tuning:
# The Tuner component launches 1 AIP Training job for flock management.
# For example, 3 workers (defined by num_parallel_trials) in the flock
# management AIP Training job, each runs Tuner.Executor.
# Then, 3 AIP Training Jobs (defined by local_training_args) are invoked
# from each worker in the flock management Job for Trial execution.
tuner = Tuner(
module_file=module_file,
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps},
tune_args=tuner_pb2.TuneArgs(
# num_parallel_trials=3 means that 3 search loops are
# running in parallel.
# Each tuner may include a distributed training job which can be
# specified in local_training_args above (e.g. 1 PS + 2 workers).
num_parallel_trials=3),
custom_config={
# Configures Cloud AI Platform-specific configs . For details, see
# https://cloud.google.com/ai-platform/training/docs/reference/rest/v1/projects.jobs#traininginput.
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
local_training_args
})
# Uses user-provided Python function that trains a model.
trainer = Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.GenericExecutor),
module_file=module_file,
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
# If Tuner is in the pipeline, Trainer can take Tuner's output
# best_hyperparameters artifact as input and utilize it in the user module
# code.
#
# If there isn't Tuner in the pipeline, either use ImporterNode to import
# a previous Tuner's output to feed to Trainer, or directly use the tuned
# hyperparameters in user module code and set hyperparameters to None
# here.
#
# Example of ImporterNode,
# hparams_importer = ImporterNode(
# instance_name='import_hparams',
# source_uri='path/to/best_hyperparameters.txt',
# artifact_type=HyperParameters)
# ...
# hyperparameters = hparams_importer.outputs['result'],
hyperparameters=(tuner.outputs['best_hyperparameters']
if enable_tuning else None),
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps},
custom_config={
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
local_training_args
})
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='species')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='SparseCategoricalAccuracy',
threshold=tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10})))
])
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
pusher = Pusher(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
custom_config=pusher_custom_config,
)
components = [
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
trainer,
model_resolver,
evaluator,
pusher,
]
if enable_tuning:
components.append(tuner)
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=components,
enable_cache=True,
beam_pipeline_args=beam_pipeline_args)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, module_file_lite: Text,
serving_model_dir: Text, serving_model_dir_lite: Text,
metadata_path: Text,
beam_pipeline_args: List[Text]) -> pipeline.Pipeline:
"""Implements the handwritten digit classification example using TFX."""
examples = external_input(data_root)
# Brings data into the pipeline.
example_gen = ImportExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(
statistics=statistics_gen.outputs['statistics'], infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
def _create_trainer(module_file, instance_name):
return Trainer(
module_file=module_file,
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
train_args=trainer_pb2.TrainArgs(num_steps=5000),
eval_args=trainer_pb2.EvalArgs(num_steps=100),
instance_name=instance_name)
# Uses user-provided Python function that trains a Keras model.
trainer = _create_trainer(module_file, 'mnist')
# Trains the same model as the one above, but converts it into a TFLite one.
trainer_lite = _create_trainer(module_file_lite, 'mnist_lite')
# TODO(b/150949276): Add resolver back once it supports two trainers.
# Uses TFMA to compute an evaluation statistics over features of a model and
# performs quality validation of a candidate model.
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='image_class')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(metrics=[
tfma.MetricConfig(
class_name='SparseCategoricalAccuracy',
threshold=tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.8})))
])
])
eval_config_lite = tfma.EvalConfig()
eval_config_lite.CopyFrom(eval_config)
# Informs the evaluator that the model is a TFLite model.
eval_config_lite.model_specs[0].model_type = 'tf_lite'
# Uses TFMA to compute the evaluation statistics over features of a model.
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
eval_config=eval_config,
instance_name='mnist')
# Uses TFMA to compute the evaluation statistics over features of a TFLite
# model.
evaluator_lite = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer_lite.outputs['model'],
eval_config=eval_config_lite,
instance_name='mnist_lite')
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)),
instance_name='mnist')
# Checks whether the TFLite model passed the validation steps and pushes the
# model to a file destination if check passed.
pusher_lite = Pusher(
model=trainer_lite.outputs['model'],
model_blessing=evaluator_lite.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir_lite)),
instance_name='mnist_lite')
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
trainer,
trainer_lite,
evaluator,
evaluator_lite,
pusher,
pusher_lite,
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=beam_pipeline_args)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text,
module_file: Text,
presto_config: presto_config_pb2.PrestoConnConfig,
query: Text, serving_model_dir: Text,
metadata_path: Text) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
# Brings data into the pipeline or otherwise joins/converts training data
example_gen = PrestoExampleGen(presto_config, query=query)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(statistics=statistics_gen.outputs['statistics'])
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Uses TFMA to compute a evaluation statistics over features of a model.
model_analyzer = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'])
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats,
transform, trainer, model_analyzer, model_validator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
)
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text,
worker_parallelism: int) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(
statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'binary_accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
model_analyzer = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(
model=trainer.outputs['model'],
model_blessing=model_analyzer.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats, transform,
trainer, model_resolver, model_analyzer, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
# LINT.IfChange
beam_pipeline_args=[
# -------------------------- Beam Args --------------------------.
'--runner=PortableRunner',
# Points to the job server started in
# setup_beam_on_{flink, spark}.sh
'--job_endpoint=localhost:8099',
'--environment_type=LOOPBACK',
'--sdk_worker_parallelism=%d' % worker_parallelism,
'--experiments=use_loopback_process_worker=True',
# Setting environment_cache_millis to practically infinity enables
# continual reuse of Beam SDK workers, improving performance.
'--environment_cache_millis=1000000',
# TODO(BEAM-7199): Obviate the need for setting pre_optimize=all. # pylint: disable=g-bad-todo
'--experiments=pre_optimize=all',
# Note; We use 100 worker threads to mitigate the issue with
# scheduling work between the Beam runner and SDK harness. Flink
# and Spark can process unlimited work items concurrently while
# SdkHarness can only process 1 work item per worker thread.
# Having 100 threads will let 100 tasks execute concurrently
# avoiding scheduling issue in most cases. In case the threads are
# exhausted, beam prints the relevant message in the log.
# TODO(BEAM-8151) Remove worker_threads=100 after we start using a # pylint: disable=g-bad-todo
# virtually unlimited thread pool by default.
'--experiments=worker_threads=100',
# ---------------------- End of Beam Args -----------------------.
# --------- Flink runner Args (ignored by Spark runner) ---------.
'--parallelism=%d' % worker_parallelism,
# TODO(FLINK-10672): Obviate setting BATCH_FORCED. # pylint: disable=g-bad-todo
'--execution_mode_for_batch=BATCH_FORCED',
# ------------------ End of Flink runner Args -------------------.
],
# LINT.ThenChange(setup/setup_beam_on_spark.sh)
# LINT.ThenChange(setup/setup_beam_on_flink.sh)
)
def _create_pipeline():
"""Implements the chicago taxi pipeline with TFX."""
examples = csv_input(_data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
infer_schema = SchemaGen(statistics=statistics_gen.outputs['statistics'])
# Performs anomaly detection based on statistics and data schema.
validate_stats = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=infer_schema.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=infer_schema.outputs['schema'],
module_file=_taxi_module_file)
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(module_file=_taxi_module_file,
examples=transform.outputs['transformed_examples'],
schema=infer_schema.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10000),
eval_args=trainer_pb2.EvalArgs(num_steps=5000))
# Uses TFMA to compute a evaluation statistics over features of a model.
model_analyzer = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
feature_slicing_spec=evaluator_pb2.FeatureSlicingSpec(specs=[
evaluator_pb2.SingleSlicingSpec(
column_for_slicing=['trip_start_hour'])
]))
# Performs quality validation of a candidate model (compared to a baseline).
model_validator = ModelValidator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'])
# This custom component serves as a bridge between pipeline and human model
# reviewers to enable review-and-push workflow in model development cycle. It
# utilizes Slack API to send message to user-defined Slack channel with model
# URI info and wait for go / no-go decision from the same Slack channel:
# * To approve the model, users need to reply the thread sent out by the bot
# started by SlackComponent with 'lgtm' or 'approve'.
# * To reject the model, users need to reply the thread sent out by the bot
# started by SlackComponent with 'decline' or 'reject'.
slack_validator = SlackComponent(
model=trainer.outputs['model'],
model_blessing=model_validator.outputs['blessing'],
slack_token=_slack_token,
slack_channel_id=_slack_channel_id,
timeout_sec=3600,
)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(model=trainer.outputs['model'],
model_blessing=slack_validator.outputs['slack_blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=_serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=_pipeline_name,
pipeline_root=_pipeline_root,
components=[
example_gen, statistics_gen, infer_schema, validate_stats,
transform, trainer, model_analyzer, model_validator,
slack_validator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
_metadata_db_root),
)
def create_e2e_components(
pipeline_root: str,
csv_input_location: str,
transform_module: str,
trainer_module: str,
) -> List[BaseComponent]:
"""Creates components for a simple Chicago Taxi TFX pipeline for testing.
Args:
pipeline_root: The root of the pipeline output.
csv_input_location: The location of the input data directory.
transform_module: The location of the transform module file.
trainer_module: The location of the trainer module file.
Returns:
A list of TFX components that constitutes an end-to-end test pipeline.
"""
example_gen = CsvExampleGen(input_base=csv_input_location)
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'])
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=transform_module)
latest_model_resolver = resolver.Resolver(
strategy_class=latest_artifact_strategy.LatestArtifactStrategy,
latest_model=Channel(type=Model)).with_id('latest_model_resolver')
trainer = Trainer(
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
base_model=latest_model_resolver.outputs['latest_model'],
transform_graph=transform.outputs['transform_graph'],
train_args=trainer_pb2.TrainArgs(num_steps=10),
eval_args=trainer_pb2.EvalArgs(num_steps=5),
module_file=trainer_module,
)
# Set the TFMA config for Model Evaluation and Validation.
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
metrics_specs=[
tfma.MetricsSpec(
metrics=[tfma.MetricConfig(class_name='ExampleCount')],
thresholds={
'accuracy':
tfma.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.5}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
])
evaluator = Evaluator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
eval_config=eval_config)
infra_validator = InfraValidator(
model=trainer.outputs['model'],
examples=example_gen.outputs['examples'],
serving_spec=infra_validator_pb2.ServingSpec(
tensorflow_serving=infra_validator_pb2.TensorFlowServing(
tags=['latest']),
kubernetes=infra_validator_pb2.KubernetesConfig()),
request_spec=infra_validator_pb2.RequestSpec(
tensorflow_serving=infra_validator_pb2.
TensorFlowServingRequestSpec()))
pusher = Pusher(
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=os.path.join(pipeline_root, 'model_serving'))))
return [
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
latest_model_resolver,
trainer,
evaluator,
infra_validator,
pusher,
]
def create_tfx_pipeline(pipeline_name: Text, pipeline_root: Text,
data_root: Text, module_file: Text,
serving_model_dir: Text, metadata_path: Text,
direct_num_workers: int) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX."""
# Brings data into the pipeline or otherwise joins/converts training data.
train_data_path = os.path.join(data_root, 'train')
test_data_path = os.path.join(data_root, 'test')
train_examples = tfrecord_input(train_data_path)
train_example_gen = ImportExampleGen(input=train_examples, instance_name='train_example_gen')
test_examples = tfrecord_input(test_data_path)
test_example_gen = ImportExampleGen(input=test_examples, instance_name='test_example_gen')
# Computes statistics over data for visualization and example validation.
train_statistics_gen = StatisticsGen(examples=train_example_gen.outputs['examples'])
# test_statistics_gen = StatisticsGen(examples=test_example_gen.outputs['examples'])
# Generates schema based on statistics files.
train_infer_schema = SchemaGen(
statistics=train_statistics_gen.outputs['statistics'],
infer_feature_shape=False)
train_transform = Transform(
examples=train_example_gen.outputs['examples'],
schema=train_infer_schema.outputs['schema'],
module_file=module_file,
instance_name='train_transformer')
test_transform = Transform(
examples=test_example_gen.outputs['examples'],
schema=train_infer_schema.outputs['schema'],
module_file=module_file,
instance_name='test_transformer')
# Uses user-provided Python function that implements a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
# need to use custom executor spec
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
transformed_examples=train_transform.outputs['transformed_examples'],
transform_graph=train_transform.outputs['transform_graph'],
schema=train_infer_schema.outputs['schema'],
train_args=trainer_pb2.TrainArgs(num_steps=20),
eval_args=trainer_pb2.EvalArgs(num_steps=10))
test_pred = custom_component.TestPredComponent(
transformed_examples=test_transform.outputs['transformed_examples'],
model=trainer.outputs['model']
)
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='Survived')],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'BinaryAccuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}))
})
])
evaluator = Evaluator(
examples=train_example_gen.outputs['examples'],
model=trainer.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
train_example_gen, train_statistics_gen, train_infer_schema, train_transform, trainer,
test_example_gen, test_transform, test_pred,
# evaluator, pusher
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
beam_pipeline_args=['--direct_num_workers=%d' % direct_num_workers])
def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
module_file: Text, serving_model_dir: Text,
metadata_path: Text,
direct_num_workers: int) -> pipeline.Pipeline:
"""Implements the Iris flowers pipeline with TFX."""
examples = external_input(data_root)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = CsvExampleGen(input=examples)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(
statistics=statistics_gen.outputs['statistics'], infer_feature_shape=True)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(
examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Uses user-provided Python function that trains a model using TF-Learn.
trainer = Trainer(
module_file=module_file,
custom_executor_spec=executor_spec.ExecutorClassSpec(GenericExecutor),
examples=transform.outputs['transformed_examples'],
transform_graph=transform.outputs['transform_graph'],
schema=schema_gen.outputs['schema'],
train_args=trainer_pb2.TrainArgs(num_steps=2000),
eval_args=trainer_pb2.EvalArgs(num_steps=5))
# Get the latest blessed model for model validation.
model_resolver = ResolverNode(
instance_name='latest_blessed_model_resolver',
resolver_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(type=ModelBlessing))
# Uses TFMA to compute an evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(label_key='variety')],
slicing_specs=[tfma.SlicingSpec()],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'sparse_categorical_accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(
examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
# Change threshold will be ignored if there is no baseline (first run).
eval_config=eval_config)
# Performs infra validation of a candidate model to prevent unservable model
# from being pushed. This config will launch a model server of the latest
# TensorFlow Serving image in a local docker engine.
infra_validator = InfraValidator(
model=trainer.outputs['model'],
examples=example_gen.outputs['examples'],
serving_spec=infra_validator_pb2.ServingSpec(
tensorflow_serving=infra_validator_pb2.TensorFlowServing(
tags=['latest']),
local_docker=infra_validator_pb2.LocalDockerConfig()),
request_spec=infra_validator_pb2.RequestSpec(
tensorflow_serving=infra_validator_pb2.TensorFlowServingRequestSpec())
)
# Checks whether the model passed the validation steps and pushes the model
# to a file destination if check passed.
pusher = Pusher(
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
infra_blessing=infra_validator.outputs['blessing'],
push_destination=pusher_pb2.PushDestination(
filesystem=pusher_pb2.PushDestination.Filesystem(
base_directory=serving_model_dir)))
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen,
statistics_gen,
schema_gen,
example_validator,
transform,
trainer,
model_resolver,
evaluator,
infra_validator,
pusher,
],
enable_cache=True,
metadata_connection_config=metadata.sqlite_metadata_connection_config(
metadata_path),
# TODO(b/142684737): The multi-processing API might change.
beam_pipeline_args=['--direct_num_workers=%d' % direct_num_workers],
)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
module_file: Text,
ai_platform_training_args: Dict[Text, Text],
ai_platform_serving_args: Dict[Text, Text],
beam_pipeline_args: List[Text],
) -> pipeline.Pipeline:
"""Implements the chicago taxi pipeline with TFX and Kubeflow Pipelines.
Args:
pipeline_name: name of the TFX pipeline being created.
pipeline_root: root directory of the pipeline. Should be a valid GCS path.
module_file: uri of the module files used in Trainer and Transform
components.
ai_platform_training_args: Args of CAIP training job. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#Job
for detailed description.
ai_platform_serving_args: Args of CAIP model deployment. Please refer to
https://cloud.google.com/ml-engine/reference/rest/v1/projects.models
for detailed description.
beam_pipeline_args: List of beam pipeline options. Please refer to
https://cloud.google.com/dataflow/docs/guides/specifying-exec-params#setting-other-cloud-dataflow-pipeline-options.
Returns:
A TFX pipeline object.
"""
# The rate at which to sample rows from the Taxi dataset using BigQuery.
# The full taxi dataset is > 200M record. In the interest of resource
# savings and time, we've set the default for this example to be much smaller.
# Feel free to crank it up and process the full dataset!
# By default it generates a 0.1% random sample.
query_sample_rate = data_types.RuntimeParameter(name='query_sample_rate',
ptype=float,
default=0.001)
# This is the upper bound of FARM_FINGERPRINT in Bigquery (ie the max value of
# signed int64).
max_int64 = '0x7FFFFFFFFFFFFFFF'
# The query that extracts the examples from BigQuery. The Chicago Taxi dataset
# used for this example is a public dataset available on Google AI Platform.
# https://console.cloud.google.com/marketplace/details/city-of-chicago-public-data/chicago-taxi-trips
query = """
SELECT
pickup_community_area,
fare,
EXTRACT(MONTH FROM trip_start_timestamp) AS trip_start_month,
EXTRACT(HOUR FROM trip_start_timestamp) AS trip_start_hour,
EXTRACT(DAYOFWEEK FROM trip_start_timestamp) AS trip_start_day,
UNIX_SECONDS(trip_start_timestamp) AS trip_start_timestamp,
pickup_latitude,
pickup_longitude,
dropoff_latitude,
dropoff_longitude,
trip_miles,
pickup_census_tract,
dropoff_census_tract,
payment_type,
company,
trip_seconds,
dropoff_community_area,
tips
FROM `bigquery-public-data.chicago_taxi_trips.taxi_trips`
WHERE (ABS(FARM_FINGERPRINT(unique_key)) / {max_int64})
< {query_sample_rate}""".format(
max_int64=max_int64, query_sample_rate=str(query_sample_rate))
# Number of epochs in training.
train_steps = data_types.RuntimeParameter(
name='train_steps',
default=10000,
ptype=int,
)
# Number of epochs in evaluation.
eval_steps = data_types.RuntimeParameter(
name='eval_steps',
default=5000,
ptype=int,
)
# Brings data into the pipeline or otherwise joins/converts training data.
example_gen = big_query_example_gen_component.BigQueryExampleGen(
query=query)
# Computes statistics over data for visualization and example validation.
statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
# Generates schema based on statistics files.
schema_gen = SchemaGen(statistics=statistics_gen.outputs['statistics'],
infer_feature_shape=False)
# Performs anomaly detection based on statistics and data schema.
example_validator = ExampleValidator(
statistics=statistics_gen.outputs['statistics'],
schema=schema_gen.outputs['schema'])
# Performs transformations and feature engineering in training and serving.
transform = Transform(examples=example_gen.outputs['examples'],
schema=schema_gen.outputs['schema'],
module_file=module_file)
# Update ai_platform_training_args if distributed training was enabled.
# Number of worker machines used in distributed training.
worker_count = data_types.RuntimeParameter(
name='worker_count',
default=2,
ptype=int,
)
# Type of worker machines used in distributed training.
worker_type = data_types.RuntimeParameter(
name='worker_type',
default='standard',
ptype=str,
)
ai_platform_training_args = copy.copy(ai_platform_training_args)
if FLAGS.distributed_training:
ai_platform_training_args.update({
# You can specify the machine types, the number of replicas for workers
# and parameter servers.
# https://cloud.google.com/ml-engine/reference/rest/v1/projects.jobs#ScaleTier
'scaleTier': 'CUSTOM',
'masterType': 'large_model',
'workerType': worker_type,
'parameterServerType': 'standard',
'workerCount': worker_count,
'parameterServerCount': 1
})
# Uses user-provided Python function that implements a model using TF-Learn
# to train a model on Google Cloud AI Platform.
trainer = Trainer(
custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_trainer_executor.Executor),
module_file=module_file,
transformed_examples=transform.outputs['transformed_examples'],
schema=schema_gen.outputs['schema'],
transform_graph=transform.outputs['transform_graph'],
train_args={'num_steps': train_steps},
eval_args={'num_steps': eval_steps},
custom_config={
ai_platform_trainer_executor.TRAINING_ARGS_KEY:
ai_platform_training_args
})
# Get the latest blessed model for model validation.
model_resolver = resolver.Resolver(
strategy_class=latest_blessed_model_resolver.
LatestBlessedModelResolver,
model=Channel(type=Model),
model_blessing=Channel(
type=ModelBlessing)).with_id('latest_blessed_model_resolver')
# Uses TFMA to compute a evaluation statistics over features of a model and
# perform quality validation of a candidate model (compared to a baseline).
eval_config = tfma.EvalConfig(
model_specs=[tfma.ModelSpec(signature_name='eval')],
slicing_specs=[
tfma.SlicingSpec(),
tfma.SlicingSpec(feature_keys=['trip_start_hour'])
],
metrics_specs=[
tfma.MetricsSpec(
thresholds={
'accuracy':
tfma.config.MetricThreshold(
value_threshold=tfma.GenericValueThreshold(
lower_bound={'value': 0.6}),
# Change threshold will be ignored if there is no
# baseline model resolved from MLMD (first run).
change_threshold=tfma.GenericChangeThreshold(
direction=tfma.MetricDirection.HIGHER_IS_BETTER,
absolute={'value': -1e-10}))
})
])
evaluator = Evaluator(examples=example_gen.outputs['examples'],
model=trainer.outputs['model'],
baseline_model=model_resolver.outputs['model'],
eval_config=eval_config)
# Checks whether the model passed the validation steps and pushes the model
# to Google Cloud AI Platform if check passed.
# TODO(b/162451308): Add pusher back to components list once AIP Prediction
# Service supports TF>=2.3.
_ = Pusher(custom_executor_spec=executor_spec.ExecutorClassSpec(
ai_platform_pusher_executor.Executor),
model=trainer.outputs['model'],
model_blessing=evaluator.outputs['blessing'],
custom_config={
ai_platform_pusher_executor.SERVING_ARGS_KEY:
ai_platform_serving_args
})
return pipeline.Pipeline(
pipeline_name=pipeline_name,
pipeline_root=pipeline_root,
components=[
example_gen, statistics_gen, schema_gen, example_validator,
transform, trainer, model_resolver, evaluator
],
beam_pipeline_args=beam_pipeline_args,
)
