TensorFlow Extended (TFX)：面向生产环境的机器学习

文 / Robert Crowe，TFX 团队

谈到机器学习，人们首先想到的是可以创建许多出色的模型，毕竟这是许多研究论文着重讨论的内容。但是，如果想将这些出色的模型运用于现实环境，就要将一个生产解决方案所需的各个方面考虑在内，比如监控、可靠性、验证等。

为了解决此种需求，Google 专门为 机器学习流水线 (Machine Learning Pipelines) 提供生产环境的支持，并由此创建了 TensorFlow Extended (TFX)。我们把这个平台分享到开源社区，这样各地的开发者皆可在生产级 TFX 流水线上创建并部署自己的模型。

注：TensorFlow Extended (TFX) 链接

https://tensorflow.google.cn/tfx

Google 之所以创建 TFX 是因为我们有此需求但没有能够满足我们需求的工具。Google 以及 Alphabet 旗下的众多公司均已在多种产品中大量使用了机器学习 。其实 TFX 并非 Google 创建的第一个机器学习流水线框架。在早期，我们曾做过一些尝试，TFX 由此发展而来，现已成为 Google 用于大多数 ML 生产解决方案的默认框架。除 Google 之外，TFX 对我们的合作伙伴也产生了深远影响，其中包括 Twitter、Airbnb 和 PayPal。

机器学习的延伸

刚开始规划在您的产品中使用 ML 时，您通常要将有关 ML 的所有方面都考虑在内。这包括为监督学习获取标记好的数据，并确保数据集全面覆盖所有可能的输入。此外，还需要在最小化特征集的维度的同时最大化其预测信息。

您需要考虑产品的公正性。您还需要考虑到一些特殊情况，尤其是在医疗保健等领域的应用中，您要对罕见但极其重要的情况作出预测（例如疑难杂症）。最后，您需要为数据的生命周期管理制定计划，这是一个动态的解决方案，随着新数据的输入和条件改变，它会不断演变。

除此之外，您需要意识到该应用将投入至生产环境。这说明您还需要满足任何生产级应用的全部需求，包括可扩展性、一致性、模块化、可测试性和安全性。您现在不仅仅是在训练模型！就本身而言，这些需求对于任何生产环境的应用部署来说都是挑战，您不能因为正在进行 ML 而将其忘在脑后。那么问题来了，如何满足这些需求，并将您的新模型投入到生产环境？

这正是 TFX 的作用所在。借助 TFX，您可以为满足生产应用部署与最佳实践的众多需求而创建一条生产级机器学习流水线。TFX 从提取数据开始，然后经过数据验证、特征工程、训练、评估和提供服务。除了 TensorFlow 本身以外，Google 还为 机器学习流水线的每个主要阶段（TensorFlow Data Validation、TensorFlow Transform 和 TensorFlow Model Analysis）创建了库。同样的，Google 还针对大规模的部署场景构建了框架，包括服务器集群 (TensorFlow Serving)、原生移动应用 (TensorFlow Lite) 和 JavaScript 应用 (TensorFlow JS)。TFX 基于这一系列库构建了的流水线组件，同时您也可以创建自己的组件。

注：TensorFlow Data Validation 链接

(https://tensorflow.google.cn/tfx/guide/tfdv)

TensorFlow Transform 链接

(https://tensorflow.google.cn/tfx/transform/get_started)

TensorFlow Model Analysis 链接
(https://tensorflow.google.cn/tfx/model_analysis/get_started)

为了将这些组件关联在一起，Google 基于流水线存储空间、配置和调度等创建了一些平面。这些平面非常重要，可用于管理并优化您的流水线以及在流水线上运行的应用。

在生产环境中部署 ML 时会遇到许多挑战，Google 不会假装能解决所有问题。这将是 ML 社区中一个不断发展的领域，我们欢迎各位建言献策。本文将简要介绍在生产环境中运用机器学习时会遇到的挑战。

什么是 “流水线” 和 “组件”？

TFX 流水线由一系列组件构成，每个组件执行一项不同的任务。组件排列成有向无环图  (DAG)。那组件又是什么？

一个 TFX 组件主要包含三个部分：Driver、Executor 和 Publisher。其中 Driver 和 Publisher 主要是模板代码，您可以修改模板代码，但一般不需要这么做。实际上，Executor 才是您插入代码并制定操作的地方。

Driver 会检查组件的状态并决定需要完成的工作，同时协调作业执行步骤并将元数据返回给 Executor。Publisher 会获取 Executor 的执行结果并更新存储的元数据，而 Executor 才是每个组件实际工作的地方。

首先，您需要使用 Python 来生成 TFX 的配置文件，这些配置文件将用于 TFX 的组件中。接下来，您需要为组件提供一些输入数据以及一个存放结果的地方（也就是元数据存储的地方）。稍后我们会详细介绍元数据存储，现在只需了解到大多数组件输入的元数据来自元数据存储，并且会将生成的结果元数据写回到元数据存储中。

因此，当数据在流水线中移动时，组件会读取前面的组件生成的元数据，并将结果元数据写入流水线中以供后面的组件使用。大多数情况下，数据会以这种方式在 TFX 流水线中流动，但也有一些例外情况，比如在流水线的开始和结束的位置。

TFX 流水线的调度

为确保所有组件井然有序，同时便于管理好流水线，需要使用到调度器 (Orchestrator) 。什么是 调度器 ？它有什么用处？

调度器是用来整合机器学习流水线、定义流水线中组件的序列并管理其执行。它提供了一个管理界面，可以用来触发任务并监控组件。

如果您仅需启动流水线的下一阶段，使用任务感知架构即可。当前一个组件完成任务后，就会立即自动启动下一个组件。任务和数据感知架构的功能非常强大，几乎是所有生产系统的标配，因为它可以存储每个组件在多次执行中产生的所有 Artifact 。获得该元数据后，可以创建更加强大的流水线，并完成许多本来很困难的任务，因此 TFX 实现了任务和数据感知流水线架构。

TFX 的开放性与可扩展性也体现在了调度器上。除了 TFX 本身的 Orchestrator 之外，Google 为 Apache Airflow 和 Kubeflow 也提供支持，您也可以根据自身需要通过编写代码来使用自定义的 Orchestrator。如果您已经有自己喜欢的其他工作流引擎，可以构建一个 Runner 与 TFX 配合使用。

元数据存储的意义

TFX 使用 ML - Metadata (MLMD) 来实现元数据存储，这是一个用于定义、存储和查询 机器学习流水线的元数据的开源库。MLMD 会将元数据存储在所关联的后端。目前的版本支持 SQLite 和 MySQL。您也可以通过编写代码来扩展 ML - Metadata 以适配绝大多数兼容 SQL 的数据库。这里有个问题，元数据中存储的究竟是什么？

首先，它会存储与已训练模型相关的信息、训练模型时使用的数据及其评估结果。我们将这类元数据称为 “Artifact”，Artifact 具有属性。数据本身存储在数据库之外，但数据的属性和位置会在元数据存储中保存。

接下来，我们会在每一个组件的每一次运行时保存它的执行记录。在机器学习流水线的漫长生命周期中，由于新数据的输入或条件的改变，机器学习流水线通常会频繁运行，保存历史记录对于调试、效果复现和审核十分重要。

最后，我们还会保存数据对象在流水线中的起源。这样，您就可以在流水线中向前或向后追踪，以了解在数据和代码发生变化时组件的运行结果又是如何变化的。这个功能对于优化或调试流水线非常重要，否则工作会变得非常困难。

元数据驱动的功能

现在，您已对元数据仓库中存储的内容有了一定的了解，下面介绍其中的部分功能。

首先，在掌握了所有数据的来源之后，您可以借此在流水线中向前或向后进行追踪。例如，查看训练模型所使用的数据、新的特征工程对评估指标所产生的影响。在一些案例中，“能够追踪数据的来源和结果” 甚至是当地规章或法律要求。

请谨记，上述方法不仅仅适用当前某一时刻的模型和结果。在吸纳新数据并重新训练模型时，您可能还想知道数据和结果在一段时间内的变化。这往往需要与昨天或上周的模型运行情况做对比，以了解结果到底是变好还是变坏，以及原因是什么。生产解决方案可不是一次性用品，根据您的需求会一直存在，可能是几个月，也可能是几年。

您还可以在必要时重新运行组件来提升流水线的效率，然后使用热启动来继续训练。注意，实际情况下经常需要处理大型数据集的，而这些数据集可能需要数小时或数天才能运行完毕。比如说您已经对模型进行了一整天的训练，并且后续还有训练要继续，这时您可以从上次离开的位置继续，而不必从头开始。只要在元数据中保存了模型的相关信息，实现这一切就会变得非常容易。

当输入的数据或代码发生改变时，您可以通过重新运行组件来提升效率。无需重新运行组件，只需从缓存中提取之前的结果即可。例如，如果重新运行只是改变了 Trainer 的参数，那么流水线依然可以重用 Artifact 的预处理数据，比如词汇表。此功能可以节省大量时间，并解决大量数据导致的数据预处理成本过高的问题。在 TFX 和 MLMD 支持的数据重用功能之外，您会看到一个更简单的 “运行流水线” 界面，无需手动选择要运行的组件。这同样也可以帮助您节省数据处理时间。总之，在元数据中保存组件的输入数据和结果将会极大地简化您的工作。

组件？哪种组件？

我们在前面介绍了调度器 (Orchestrator)，下面我们将介绍 TFX 里面的标准组件。

谈到组件就绕不开 Apache Beam

在介绍标准组件之前，我们先要介绍 Apache Beam。要进行大数据分布式处理，特别是面对机器学习这样的数据密集型计算，需要支持分布式处理的流水线框架，比如 Apache Spark、Apache Flink 或 Google Cloud Dataflow。

而 Apache Beam 是一个统一编程模型，在上述的多个执行引擎上通用。大多数的 TFX 组件运行在 Apache Beam 之上。采用 Beam，您可以直接使用现有的分布式处理框架，也可以选择自己喜欢的框架，而无需强制使用我们提供的框架。避免了切换框架的烦恼。目前，Beam Python 可以在 Flink、Spark 和 Dataflow Runner 运行器上运行，日后会新增更多的 Runner。Beam 中还有一个 Direct Runner，支持在本地系统中（比如自己的笔记本电脑）在开发环境下运行 TFX 流水线。

标准组件

安装 TFX 时，包含了一套相当完整的标准组件，每一个组件都是针对生产级 机器学习流水线的一个部分而设计的。例如，Transform 组件使用 Apache Beam 来执行特征工程转换，如创建词汇表或执行主成分分析 (PCA)。您可以在 Flink 或 Spark 集群上运行这些转换，亦可在使用  Google Cloud 上的 Dataflow 完成此操作。Apache Beam 具有可移植性，您无需更改代码，即可在平台之间来回迁移。

Trainer 组件实际只使用 TensorFlow。还记得满脑子想的都是训练模型的情形吗？这里的代码就是起到这个作用的。请注意，TFX 目前仅支持 tf.estimators。这里列出了有关兼容性的详细信息。

有些组件十分简单。例如，Pusher 组件仅需 Python 即可完成工作。

当您将这些组件整合到一起，并使用调度器进行管理时，您就拥有了一个 TFX 流水线。您可在一端接收数据，在另一端将 SavedModel 推送到一个或多个部署目标中。这些部署目标包括模型存储区，如 TensorFlow Hub、使用 TensorFlow JS 的 JavaScript 环境、使用 TensorFlow Lite 的原生移动应用、使用 TensorFlow Serving 的应用系统，或上述全部。

下面我们详细地介绍各个组件。

读取数据

首先，使用 ExampleGen 提取输入数据。ExampleGen 是在 Beam 上运行的组件之一。该组件会从各类支持的来源和类型中读取数据，同时将其划分为训练数据与评估数据，并将其格式转换为 tf.examples。ExampleGen 的配置过程十分简单，仅需两行 Python 代码。

接着，StatisticsGen 会使用 Beam 对数据进行遍历（即一个完整的 epoch），并为每个特征进行描述统计。此时，StatisticsGen 会利用 TensorFlow Data Validation (TFDV) 库。TFDV 库支持一些可视化工具，您可以在 Jupyter Notebook 中运行这些可视化工具。您便能借此探索和理解自己的数据，并发现可能存在的问题。这是典型的数据处理流程。在训练模型之前，我们都要完成这个环节。

下一个组件 SchemaGen 也使用 TensorFlow Data Validation 库。该组件对 StatisticsGen 处理过的统计数据，并尝试推断特征的基本属性，包括特征值的数据类型、值的范围和类别。可以根据需要检查和调整结构（Schema），例如为值添加您想要的新类别。

接下来，组件 ExampleValidator 会获取 StatisticsGen 的统计数据和结构（可以是 SchemaGen 直接输出的结果或用户调整后的结果），并查找问题。该组件会查找不同类别的异常，包括丢失的值或与模式不匹配的值、训练－应用偏差、数据偏移，并生成异常报告。请谨记，您的流水线一直在接收新数据，因此需要第一时间发现问题。

特征工程

Transform 是较为复杂的一个组件，需要更多的配置和额外的代码。Transform 使用 Beam 完成特征工程，通过转换特征来提高模型的性能。例如，Transform 可以创建词汇表、对值进行分桶或基于您的输入运行 PCA。编写什么代码取决于您为模型和数据集进行什么特征工程。

Transform 会对您的数据进行完整的遍历（即一个完整的 epoch），并输出两类结果。比如，要计算一个特征的中间值或标准差，在这类情况下，所有示例的数字均相同，Transform 就会输出一个常量。在诸如对值进行标准化的情况下，不同示例的值会有不同的结果，因此 Transform 会输出 TensorFlow 算子。

然后，Transform 会输出包含这些常量和算子的 TensorFlow 计算图。该图为封闭图，因此其包含使用这些转换所需的全部信息，并且会形成模型的输入阶段。这意味着，训练与应用之间通过使用相同的转换保持一致，而这可以消除训练 / 应用偏差。与此相反的情形是：您想把模型从训练环境移至实际应用环境或应用程序内，并在两处环境使用相同的特征工程，您希望出来的效果是相同的，但有时却发现并非如此。我们将此情况称为训练 / 应用偏差，而 Transform 会在您运行模型的任何地方使用完全相同的代码，从而消除该偏差。

训练模型

现在，您终于准备好去训练模型了，这个环节也是多数人在谈到机器学习的时候直接想到的事情。Trainer 会接收 Transform 的转换图和数据，以及 SchemaGen 的结构，并使用您的建模代码训练模型。模型训练和以往别无二致，但在训练完成时，Trainer 会保存两个不同的 SavedModel。一个是 SavedModel，用于部署至生产环境；而另一个为 EvalSavedModel，用于分析模型性能。

您可以为 Trainer 自定义配置，例如步数，以及是否使用 “热启动” 方式。您为 Trainer 创建的代码为建模代码，因此可以根据需要决定它的复杂程度。

如要监控和分析训练过程，您可以如往常一样使用 TensorBoard。在此界面下，您可以查看当前的模型训练运行情况，或比较多个模型训练运行的结果。需要注意的是，只有配置了 ML-Metadata 元数据存储（如上文所述），才能支持这些功能。使用 TFX，您可以很轻松地完成此类比较，而这通常具有启发意义。

尚不满意？

现在您已训练完模型，结果如何？Evaluator 组件会接收 Trainer 创建的 EvalSavedModel 和原始输入数据，并使用 Beam 和 TensorFlow 模型分析库进行深入研究。这不仅是查看整个数据集中的顶层结果，而是更深入地研究数据集的各个部分。这很重要，因为在使用模型期间，每位用户的体验都取决于其个人数据点。您的模型可能在您的整个数据集上表现优异，但如果在用户提供的数据点上表现糟糕，则用户体验便会很差。我们会在日后的文章中进一步讨论此问题。

在了解模型的性能后，是否应将该模型投入生产？相较于生产环境内的现有模型，该模型表现更优还是更差？您可能并不想在没有对比的情况下强推新模型。因此，ModelValidator 组件会使用 Beam 进行这类比较，通过您定义的标准来决定是否将新模型投入生产。

如果 ModelValidator 确定您的新模型已经可以投入生产，则 Pusher 会完成将其部署到目标环境的工作。这些目标可能是 TensorFlow Lite（若正在开发移动应用）、TensorFlow JS（若要部署至 JavaScript 环境）、TensorFlow Serving（若要部署至应用系统）或以上所有环境。

注：TensorFlow Lite 链接

(https://tensorflow.google.cn/lite)

TensorFlow JS 链接

(https://tensorflow.google.cn/js)

TensorFlow Serving 链接

(https://tensorflow.google.cn/tfx/guide/serving)

未来目标

本文旨在概述 TFX 和 机器学习流水线的基本知识，并介绍其主要概念。在接下来的文章中，我们将深入探讨 TFX，包括讨论如何根据所需来扩展 TFX。TFX 是开源项目，因此 我们 也鼓励开源应用和机器学习社区帮助我们改进该项目。您可以浏览 TFX 开发者教程，轻松开启入门学习！

注：TFX 开发者教程 链接

(https://tensorflow.google.cn/tfx/tutorials)

致谢

感谢 Robert Crowe、 Konstantinos (Gus) Katsiapis 和 Kevin Haas 及 TFX 团队。

如果您想详细了解 TensorFlow 的相关内容，请参阅以下文档。这些文档深入探讨了这篇文章中提及的许多主题：

• 观看优酷上的相关视频

   (https://v.youku.com/v_show/id_XNDE3NTQxMzYwMA==.html?spm=a2hzp.8253876.0.0&f=52238552)

• TFX：基于 TF 的生产级机器学习平台 (KDD 2017)

  (https://www.kdd.org/kdd2017/papers/view/tfx-a-tensorflow-based-production-scale-machine-learning-platform)
  

• TensorFlow

  (https://tensorflow.google.cn)
  

• TensorFlow Extended (TFX)

  (https://tensorflow.google.cn/tfx)

• TFX 开发者教程

  (https://tensorflow.google.cn/tfx/tutorials

• 联系我们
  (https://github.com/tensorflow/tfx/blob/master/CONTRIBUTING.md)

点击 “阅读原文”，立即观看 系列视频