TFX: A TensorFlow-Based Production-Scale Machine Learning Platform

By Denis Baylor, Eric Breck, Heng-Tze Cheng, Noah Fiedel, Chuan Yu Foo, Zakaria Haque, S. Haykal, M. Ispir, Vihan Jain, L. Koc and 12 more. 2017

478 citations recorded by the source corpus.

Creating and maintaining a platform for reliably producing and deploying machine learning models requires careful orchestration of many components---a learner for generating models based on training data, modules for analyzing and validating both data as well as models, and finally infrastructure for serving models in production. This becomes particularly challenging when data changes over time and fresh models need to be produced continuously. Unfortunately, such orchestration is often done ad hoc using glue code and custom scripts developed by individual teams for specific use cases, leading to duplicated effort and fragile systems with high technical debt. We present TensorFlow Extended (TFX), a TensorFlow-based general-purpose machine learning platform implemented at Google. By integrating the aforementioned components into one platform, we were able to standardize the components, simplify the platform configuration, and reduce the time to production from the order of months to weeks, while providing platform stability that minimizes disruptions. We present the case study of one deployment of TFX in the Google Play app store, where the machine learning models are refreshed continuously as new data arrive. Deploying TFX led to reduced custom code, faster experiment cycles, and a 2% increase in app installs resulting from improved data and model analysis.

Subfields: Artificial Intelligence

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Mohit Bansal
2. Yang Liu
3. Percy Liang
4. Quoc V. Le
5. Sergey Levine
6. Graham Neubig
7. Yoav Goldberg
8. Li Fei-Fei

Citations in the corpus

• AlphaClean: Automatic Generation of Data Cleaning Pipelines (2019, methodology): “An initial architecture is to directly apply recent hyperparameter tuning approaches for machine learning pipelines and neural network model search [7, 20, 32, 33, 44].”
• A Survey on Data Collection for Machine Learning: A Big Data - AI Integration Perspective (2018, methodology): “This problem occurs frequently in practice, so production machine learning platforms like TensorFlow Extended (TFX) [138] have separate components [139] to reduce data errors as much as possible though analysis and validation.”
• The ML test score: A rubric for ML production readiness and technical debt reduction (2017, background): “TFX system[18] can be beneficial here as it then allows building a generic integration test.”
• Recommending what video to watch next: a multitask ranking system (2019, methodology): “Specifcally, we use Tensor Processing Units (TPUs) to train our model and serve it using TFX Servo [4] 3.”
• Software Engineering for Machine Learning: A Case Study (2019, background): “In the last five years, there have been multiple efforts in industry to automate this process by building frameworks and environments to support the ML workflow and its experimental nature [1], [22], [23].”
• Automating Large-Scale Data Quality Verification (2018, background): “As a consequence, modern machine learning platforms begin to have explicit data validation components [7, 5, 9].”
• Accelerating the Machine Learning Lifecycle with MLflow (2018, methodology): “In practice, an organization will need to run models from multiple ML libraries, TensorFlow versions, etc., and has to design its own infrastructure for this task.”
• Deep Neural Networks and Tabular Data: A Survey (2021, background): “This might be an issue for high-load systems [74], in the case of categorical variables with high cardinality and growing data set size.”
• The What-If Tool: Interactive Probing of Machine Learning Models (2019, methodology): “Out-of-the-box support, with no need for any custom coding by the user in alignment with user need N1, is provided for those making use of the TensorFlow Extended (TFX) pipeline [22], through WIT’s support of TFRecord files, TensorFlow Estimator models, and model serving through TensorFlow Serving.”
• Simple, Distributed, and Accelerated Probabilistic Programming (2018, background): “system: this makes them difficult to compose with arbitrary computation, and with the broader machine learning ecosystem, such as production platforms [5].”
• Software Engineering for AI-Based Systems: A Survey (2021, background): “Another large category is frameworks or platforms (11 studies), ranging from a generic level of abstraction for end-to-end ML application development and deployment [14, 41, 83, 99, 113] to available tool support for managing reliable ML applications [18], packaging and sharing ML models as reusable”
• Data collection and quality challenges in deep learning: a data-centric AI perspective (2021, background): “Schema-based validation [27, 122] is widely used in practice.”
• Machine Learning Testing: Survey, Landscapes and Horizons (2019, background): “1), test input or test oracle design [116], [120], [121], [123], [130], [154], [162], and test tool design [165], [167], [182] (see more in Section 5).”
• Rafiki: Machine Learning as an Analytics Service System (2018, methodology): “TensorFlow serving [2] and TFX [21] provide inference service for models trained using Tensorflow.”
• Data-centric Artificial Intelligence: A Survey (2023, unknown): “For instance, in Reference [17], a data validation module trains a machine learning model on a training set with expected data schema and generalizes it to identify potential problems in unseen scenarios.”
• An anatomy for neural search engines (2019, background): “This can lead to a lot of glue code, which is hard to maintain and evolve, and require a lot of re-work [Baylor et al., 2017].”
• Towards Accountability for Machine Learning Datasets: Practices from Software Engineering and Infrastructure (2020, unknown): “than evaluating model performance metrics in isolation [23]. In these contexts, the model itself will form only one part of the deployed system, requiring infrastructure for deployment and monitoring [14]. Dataset development documentation outlined above plays a critical role in defining the contrac”