Package: tfx-helper
tfx-helperA helper library for TFXWhy?This package contains small utilities that help in creation ofTFXpipelines:supports running pipeline locally or onVertex AI Pipelines(with all training, tuning and serving also happening insideVertex AI),abstracts local vs cloud execution environment outside of the pipeline definition (no need forif use_gcp:conditions inside your pipeline code - write a uniform pipeline creation code and run it on both local and cloud),construct complexcustom_configforTFXextension components for you (configuration of extension components is complex and not well documented - we did the research for you and are exposing a simple API),enable GPU in training/tuning/serving with a single argument,enable declaring per-component resource requirements (you can now runEvaluatorcomponent on a beefier machine if you have a large model),use generator function syntax in pipeline definition to avoid boilerplate,avoid passing a hundred parameters into your pipeline definition (cloud configuration, likeservice_accountis now only part of cloud-targeted runner)In addition we provide a set of helper functions for visualization of TFX pipeline artifacts.How?Installpipinstalltfx-helperPipeline definitionUse our helper component interface in you pipeline definition.Return a collection of components.For multi-version components (Trainer,Tuner,Pusher) construction use the helper.fromtfx_helper.interfaceimportPipelineHelperInterfacedefcreate_pipeline(pipeline_helper:PipelineHelperInterface,# pass in the helper as interface*,# all your pipeline parameterstrain_epochs:int,# maximum number of training epochs in trainer...# other parameters)->Iterable[BaseComponent]:# return a collection of components...# create `Transform` in the usual waytransform=tfx.components.Transform(examples=example_gen.outputs["examples"],schema=schema_gen.outputs["schema"],preprocessing_fn="models.preprocessing.preprocessing_fn",splits_config=tfx.proto.SplitsConfig(analyze=["train","valid"],transform=["train","valid","eval"],),)yieldtransform...# use the helper to create a `Trainer` in a uniform waytrainer=pipeline_helper.construct_trainer(run_fn="models.model.run_fn",examples=transform.outputs["transformed_examples"],transform_graph=transform.outputs["transform_graph"],schema=schema_gen.outputs["schema"],hyperparameters=hparams,train_args=tfx.proto.TrainArgs(splits=["train"]),eval_args=tfx.proto.EvalArgs(splits=["valid"]),# custom parameters to the training callbackcustom_config={"epochs":train_epochs,"patience":train_patience},)yieldtrainer...Pipeline local runnerCreate a pipeline runner that will take your uniform pipeline definition and materialize
it for running locally (throughDirectRunner):fromtfx_helper.localimportLocalPipelineHelperdefrun()->None:"""Create and run a pipeline locally."""input_dir=...output_dir=...serving_model_dir=...# Create pipeline helper instance of local flavour.pipeline_helper=LocalPipelineHelper(pipeline_name="sentimentanalysis",output_dir=output_dir,# Where should the model be pushed tomodel_push_destination=tfx.proto.PushDestination(filesystem=tfx.proto.PushDestination.Filesystem(base_directory=serving_model_dir)),)components=create_pipeline(# Pass our pipeline helper instancepipeline_helper,# The rest of the parameters are pipeline-specific.data_path=input_dir,...# other arguments)# Run the pipelinepipeline_helper.create_and_run_pipeline(components)Notice that no cloud-specific configuration was needed neither in the runner nor in the
pipeline definition.Pipeline cloud runnerCreate a pipline runner that will take you uniform pipeline definition and materialize
it for running in the cloud (onVertex AI PipelinesthroughKubeflowV2DagRunner):fromtfx_helper.interfaceimportResourcesfromtfx_helper.vertex_aiimportVertexAIPipelineHelperdefrun()->None:output_dir="gs://..."# minimal (less than the standard `e2-standard-4`) resource for components# that won't execute computationsminimal_resources=Resources(cpu=1,memory=4)# create a helper instance of cloud flavourpipeline_helper=VertexAIPipelineHelper(pipeline_name="...",output_dir=output_dir,google_cloud_project="...",google_cloud_region="europe-west4",# all the components will use our custom image for runningdocker_image="europe-west4-docker.pkg.dev/.../...-repo/...-image:latest",service_account="...@....iam.gserviceaccount.com",# name of the Vertex AI Endpointserving_endpoint_name="...",# Number of parallel hyperparameter tuning trailsnum_parallel_trials=2,# GPU for Trainer and Tuner componentstrainer_accelerator_type="NVIDIA_TESLA_T4",# Machine type for Trainer and Tuner componentstrainer_machine_type="n1-standard-4",# GPU for serving endpointserving_accelerator_type="NVIDIA_TESLA_T4",# Machine type for serving endpointserving_machine_type="n1-standard-4",# Override resource requirements of components. The dictionary key is the ID# of the component (usually class name, unless changed with `with_id` method).resource_overrides={# evaluator needs more RAM than standard machine can provide"Evaluator":Resources(cpu=16,memory=32),# training is done as Vertex job on a separate machine"Trainer":minimal_resources,# tuning is done as Vertex job on a separate set of machines"Tuner":minimal_resources,# pusher is just submitting a job"Pusher":minimal_resources,},)# Run the pipelinecomponents=create_pipeline(pipeline_helper,# Input data in Cloud Storagedata_path="gs://...",...# other arguments)# Run the pipelinepipeline_helper.create_and_run_pipeline(components,enable_cache=True)Viewing pipeline artifactsThe TFX pipeline artifacts can be visualized inside Jupyter Notebook cells using
(mostly) custom interactive widgets. In order to be able to display those widgets,
the following prerequisites need to be fulfilled:widget extensions need to be installed,jupyternbextensionenable--pywidgetsnbextension--sys-prefix
jupyternbextensioninstall--py--symlinktensorflow_model_analysis--sys-prefix
jupyternbextensionenable--pytensorflow_model_analysis--sys-prefixnotebook needs to be trusted.Getting recent artifact directoryEach TFX artifact is stored in the TFX pipeline's output location.
Depending on whether you run your pipeline locally or on VertexAI the structure
of the directories is different. We provide helpers that can assist you in retrieving
the newest artifact locations. In local environment the helper uses directory traversal
to find the correct path. With VertexAI the helper communicates with VertexAI ML Metadata
Service.For local pipeline runs:fromtfx_helper.artifact_finder.localimportNewestLocalPathGetterfromtfx_helper.artifact_finder.interfaceimportArtifactPathGetterInterfacepath_getter:ArtifactPathGetterInterface=NewestLocalPathGetter(artifact_dir=pipeline_output,pipeline_name=pipeline_name)best_hparams_path=path_getter('Tuner','best_hyperparameters')evaluation_path=path_getter('Evaluator','evaluation')For VertexAI pipeline runs:fromtfx_helper.artifact_finder.interfaceimportArtifactPathGetterInterfacefromtfx_helper.artifact_finder.vertex_aiimportNewsetRunVertexAIPathGetterpath_getter:ArtifactPathGetterInterface=NewsetRunVertexAIPathGetter(pipeline_name=pipeline_name,region=gcp_region,project=gcp_project)best_hparams_path=path_getter('Tuner','best_hyperparameters')evaluation_path=path_getter('Evaluator','evaluation')The obtained directory can be used with our set of helper functions that help you
preview the artifacts in an easy manner.Please note that the name of a pipeline component can be customized. If you define
in your pipeline:stats_gen=tfx.components.StatisticsGen(...)then you should use:evaluation_path=path_getter('StatisticsGen','statistics')but if you customize the name:stats_gen = tfx.components.StatisticsGen(...).with_id("raw_stats_gen")then you should use:evaluation_path=path_getter('raw_stats_gen','statistics')Viewing dataset statisticsUse thedisplay_statshelper to compare statistics of two splits of the same dataset:fromtfx_helper.visualization.display_statsimportdisplay_statsdisplay_stats(path_getter('raw_stats_gen','statistics'))Use thecompare_statshelper to compare statistics between two datasets.fromtfx_helper.visualization.display_statsimportcompare_statscompare_stats(left_dir=path_getter('raw_stats_gen','statistics'),right_dir=path_getter('inference_stats_gen','statistics'))Viewing dataset schemaUse thedisplay_schemahelper to preview dataset schema:fromtfx_helper.visualization.display_schemaimportdisplay_schemadisplay_schema(path_getter('SchemaGen','schema'))Viewing dataset anomaliesIf you useExampleValidatorcomponent, you can preview its anomalies detection report
usingdisplay_anomalieshelper:fromtfx_helper.visualization.display_anomaliesimportdisplay_anomaliesdisplay_anomalies(path_getter('ExampleValidator','anomalies'),split_name='all')Viewing best hyper parametersTo view the values of best hyperparameters chosen by theTunercomponent, you can use
thedisplay_hyperparamshelper:fromtfx_helper.visualization.display_hyperparamsimportdisplay_hyperparamsdisplay_hyperparams(path_getter('Tuner','best_hyperparameters'))Viewing model evaluation resultsIf you useslicing_specsinEvalConfigto yourEvaluatorcomponent, then you might
be willing to view what kind of slices were produced during evaluation:fromtfx_helper.visualization.display_metricsimportget_slice_namesget_slice_names(path_getter('Evaluator','evaluation'))To view the overall metrics for your model usedisplay_metrics:fromtfx_helper.visualization.display_metricsimportdisplay_metricsdisplay_metrics(path_getter('Evaluator','evaluation'))To view metrics by slice use:display_metrics(path_getter('Evaluator','evaluation'),column='HomePlanet')To view the overall plots for your model usedisplay_plots:fromtfx_helper.visualization.display_metricsimportdisplay_plotsdisplay_plots(path_getter('Evaluator','evaluation'))To view plots for a specific slice use:display_plots(path_getter('Evaluator','evaluation'),slice_key={'HomePlanet':'Earth'})To check whether the candidate model passed validation (when you have thresholds configured inEvalConfig)
usepassed_validation:fromtfx_helper.visualization.display_metricsimportpassed_validationpassed_validation(path_getter('Evaluator','evaluation'))Viewing binary classification confusion matricesTo view a binary classification confusion matrix constructed from the evaluation result
of your model (data availability determined by presence ofConfusionMatrixPlotinEvalConfigto yourEvaluatorcomponent):fromtfx_helper.visualization.confusion_matriximportplot_binary_classification_confusion_matrixplot_binary_classification_confusion_matrix(path_getter('Evaluator','evaluation'))We provide a pipeline component for finding the best (by geometric mean) threshold
to use from the thresholds gathered during evaluation:fromtfx_helper.components.threshold_optimizer.componentimport(BinaryClassificationThresholdOptimizer,)threshold_optimizer=BinaryClassificationThresholdOptimizer(model_evaluation=evaluator.outputs["evaluation"])To view the confusion plot with the optimized threshold use:fromtfx_helper.visualization.threshold_optimizationimportload_best_thresholdbest_threshold=load_best_threshold(path_getter('ThresholdOptimizer','best_threshold'))plot_binary_classification_confusion_matrix(path_getter('Evaluator','evaluation'),threshold=best_threshold)More infoLink to article describing creation ofTFXpipeline for sentiment analysis using
this helper library:https://dlabs.ai/resources/courses/bert-sentiment-analysis-on-vertex-ai-using-tfx/