def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_root: Text,
module_file: Text,
ai_platform_training_args: Dict[Text, Text],
ai_platform_serving_args: Dict[Text, Text],
enable_tuning: bool,
beam_pipeline_args: Optional[List[Text]] = None) -> pipeline.Pipeline:
"""Implements the Iris flowers 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 Iris flowers 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.
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.
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.
"""
examples = external_input(data_root)
# 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=n1-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=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)
# 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 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)
pusher = 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
})
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,
serving_model_uri: Text,
data_root_uri: Union[Text, data_types.RuntimeParameter],
schema_folder_uri: Union[Text, data_types.RuntimeParameter],
train_steps: Union[int, data_types.RuntimeParameter],
eval_steps: Union[int, data_types.RuntimeParameter],
beam_pipeline_args: List[Text],
trainer_custom_executor_spec: Optional[executor_spec.ExecutorSpec] = None,
trainer_custom_config: Optional[Dict[Text, Any]] = None,
enable_tuning: Optional[bool] = False,
enable_cache: Optional[bool] = False,
metadata_connection_config: Optional[
metadata_store_pb2.ConnectionConfig] = None
) -> pipeline.Pipeline:
"""Trains and deploys the Keras Covertype Classifier with TFX and AI Platform Pipelines."""
# 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,
source_uri=schema_folder_uri,
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, 3 workers (defined by num_parallel_trials) in the flock
# management AI Platform Training job, each runs Tuner.Executor.
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=3 means that 3 search loops are running in parallel.
num_parallel_trials=3),
custom_config=custom_config)
# Trains the model using a user provided trainer function.
trainer = Trainer(
custom_executor_spec=trainer_custom_executor_spec,
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=trainer_custom_config)
# 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)
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_uri)))
components = [
examplegen, statisticsgen, schemagen, import_schema, examplevalidator,
transform, trainer, resolver, evaluator, 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,
metadata_connection_config=metadata_connection_config)
def create_pipeline(
pipeline_name: Text,
pipeline_root: Text,
data_root: Text,
module_file: Text,
ai_platform_training_args: Dict[Text, Text],
ai_platform_serving_args: Dict[Text, Text],
enable_tuning: bool,
beam_pipeline_args: List[Text],
) -> 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.
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.
enable_tuning: If True, the hyperparameter tuning through CloudTuner is
enabled.
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.
"""
# 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)
# 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 of
# parallel tuning. For example, 2 workers (defined by num_parallel_trials)
# in the flock management AIP Training job, each runs a search loop for
# trials as shown below.
# Tuner component -> CAIP job X -> CloudTunerA -> tuning trials
# -> CloudTunerB -> tuning trials
#
# Distributed training for each trial depends on the Tuner
# (kerastuner.BaseTuner) setup in tuner_fn. Currently CloudTuner is single
# worker training per trial. DistributingCloudTuner (a subclass of
# CloudTuner) launches remote distributed training job per trial.
#
# E.g., single worker training per trial
# ... -> CloudTunerA -> single worker training
# -> CloudTunerB -> single worker training
# vs distributed training per trial
# ... -> DistributingCloudTunerA -> CAIP job Y -> master,worker1,2,3
# -> DistributingCloudTunerB -> CAIP job Z -> master,worker1,2,3
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.
num_parallel_trials=3),
custom_config={
# Note that this TUNING_ARGS_KEY will be used to start the CAIP job
# for parallel tuning (CAIP job X above).
#
# num_parallel_trials will be used to fill/overwrite the
# workerCount specified by TUNING_ARGS_KEY:
# num_parallel_trials = workerCount + 1 (for master)
ai_platform_tuner_executor.TUNING_ARGS_KEY:
ai_platform_training_args,
# This working directory has to be a valid GCS path and will be used
# to launch remote training job per trial.
ai_platform_tuner_executor.REMOTE_TRIALS_WORKING_DIR_KEY:
os.path.join(_pipeline_root, 'trials'),
})
# 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:
ai_platform_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 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)
pusher = 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
},
)
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_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 = 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, output_config=output)
# 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, 3 workers (defined by num_parallel_trials) in the flock
# management AI Platform Training job, each runs Tuner.Executor.
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=3 means that 3 search loops are running in parallel.
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:
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)
