了解 Waze 如何利用 TFX 扩大可以投入生产的 ML 规模

发布人：Waze 数据科学家 Gal Moran、Iris Shmuel 和 Daniel Marcous

Waze

Waze 是世界上最大的基于社区的交通和导航应用。该应用借助实时数据来帮助用户避开路上的障碍，享受无忧旅途。除了移动导航功能，Waze 还提供网络平台、拼车、合作服务、广告投放等功能。如此广泛的服务组合带来了多种技术挑战和丰富案例。

 Waze 的 ML 应用

Waze 依赖许多 ML 解决方案，其中包括：

• 预测 ETA

• 匹配乘客和司机（拼车）

• 推送适当的广告

但是，要把这些事情做好并达到“生产级别”并不是那么容易。这类项目通常需要复杂的相关基础架构，以便将其投入生产，因此需要多位工程师（数据科学家、软件工程师和软件可靠性工程师）一起投入大量时间。如果将大规模数据、低延迟（实际上为实时）推理、多样性案例和大量地理空间数据等 Waze 的特殊要求结合在一起，则难度更甚。

• 复杂的相关基础架构
  

  https://papers.nips.cc/paper/5656-hidden-technical-debt-in-machine-learning-systems.pdf
  

上述理由很好地解释了，为什么不经周密安排就开始在 Waze 中实施 ML 会造成混乱局面。对我们来说，这表现为：

• 使用多个 ML 框架——你能想到的都有（sklearn、XGBoost、TensorFlow、fbprophet、Java PMML、人工设定的框架等等）

• ML 和运算脱节——模型和特征工程被工程师嵌入 (Java) 后端服务器，监控和验证能力有限

• 培训、验证和部署使用半手动操作

• 从想法到投入生产的开发周期异常漫长

总的来说，数据科学家最终把大量的时间耗费在运算和监控上，而不能专注于实际的建模和数据处理。

业务发展到一定程度时，我们决定整改这混乱的局面，在自动化和流程方面投资，以便能够更快地扩大业务规模。我们决定采用全周期数据科学理念，重点投资能够大幅提高速度和质量的方法。这意味着，在我们想要建立的这个新方法中，只需一个数据科学家就能完成从研究到生产级服务的产品周期。

在新方法中，数据科学家可以直接为生产做出贡献，发挥自己最大优势。他们可专注于建模和数据处理，获得许多基础架构和开箱即用的运算。虽然上述愿景还没有完全实现，但我们觉得本文所阐述的内容对实现最终目的来说至关重要。

Waze 的 ML 堆栈

为将上述理念转化为技术规范，我们决心创建一个简单、稳定、自动化和统一的方法来构建 Waze 的 ML 流水线。

深入研究技术要求后，我们得出了以下评判标准：

• 简单——易于理解、使用和操作

• 可托管——没有服务器，没有硬件，只有代码

• 可自定义——免费获得简单的东西，但也有足够的灵活性，可以为 5% 需要跳出局限的东西而激发奇思妙想

• 可扩容——自动可扩容的数据处理、训练、推理

• 利用 Python——我们需要一些已经可以投入生产的工具，可以兼容当今大多数工具和代码，并且适合一般的数据科学家。现在除了 Python，几乎没有其他选择

基于上述原因，我们选择 TFX 及其内置组件来提供这些功能，大部分是开箱即用的。

• TFX
  

  https://tensorflow.google.cn/tfx

值得一提的是，Waze 是在 Google Cloud Platform (GCP) 上运行其技术栈。

GCP 恰好可提供一套名为 Vertex AI 的工具。Waze 就是在此 ML 基础架构平台上构建的。虽然我们使用了 Vertex AI 托管式服务的许多组件，但我们将在本文中重点介绍 Vertex Pipelines，这是 ML 流水线的一个框架，能够帮助我们封装 TFX（或任何流水线）的复杂性和设置。

• Vertex Pipelines
  

  https://cloud.google.com/vertex-ai/docs/pipelines

连同我们的数据技术栈，Waze 的整体 ML 架构（均符合可托管、可扩容，且利用 Python 等标准）如下：

细心的读者会在这里注意到所谓的警告：我们所有流程都通过 TensorFlow 进行。

TFX 指的是 TensorFlow（尽管这种说法已经不再完全准确了，但我们假设它是这个意思）。

• 已经不再完全准确了
  

  https://github.com/tensorflow/tfx/blob/master/tfx/examples/penguin/experimental/penguin_utils_sklearn.py

当你处理许多不同的案例时，这些警告一开始可能看起来会有点吓人。

幸运的是，TF 生态系统很丰富，而且 Waze 的优点是有规模足够大的数据，神经网络会收敛。

自从开始这项工作以来，我们还没有发现在哪个案例中， TF 不能像其他框架那样更好地或充分地解决问题（在这里不讨论微观的百分点，不进行 Kaggle 比赛，而是想得到一些能够投入生产的东西）。

Waze TFX

您可能认为选择 TFX 和 Vertex 流水线已解决我们所有的问题，但这并不完全正确。

为了让这项工作真正变得简单，我们不得不写一些“胶水代码”（整合上述架构图中的各种产品），并提取足够的细节，以便普通的数据科学家能够有效和快速地使用这些东西。

这能带来以下好处：

• 可以淘汰样板文件

• 可隐藏所有通用的 TFX 组件，这让数据科学家可以只专注于特征工程和建模，并免费获得整个流水线

• 可生成基于 BigQuery 的训练/评估分块

• 可提供预先实现的可选通用特征转换（例如，扩容、归一化、插补）

• 可提供预先实现的 Keras 模型（例如 DNN/RNN 模型。其类似 TF Estimator，但在具有 TFX 特征的 Keras 中）

• 效用函数（例如，TF 列的准备）

• tf.transform 特征工程代码的单元测试框架

• 使用安装了所有 TFX 软件包的云运行实例，从 Airflow 中协调和调度流水线运行（无需在 Airflow Composer 上安装）

我们已经将其全部放入一个易于使用的 Python 软件包中，称为 “wze-data-tfx”。

在上文中，我们为数据科学家提供了一个超级详细的演示、使用指南和代码模板，所以常见的 DS 工作流是：创建分支，改变配置，稍微调整一下代码，部署。

作为参考，这里提供了一个简单的 waze-data-tfx 流水线：

1.配置

_DATASET_NAME = 'tfx_examples'
_TABLE_NAME = 'simple_template_data'

_LABEL_KEY = 'label'
_CATEGORICAL_INT_FEATURES = {
  "categorical_calculated": 2,
}
_DENSE_FLOAT_FEATURE_KEYS = ["numeric_feature1", "numeric_feature2"]
_BUCKET_FEATURES = {
  "numeric_feature1": 5,
}
_VOCAB_FEATURES = {
  "categorical_feature": {
      'top_k': 5,
      'num_oov_buckets': 3
  }
}

_TRAIN_BATCH_SIZE = 128
_EVAL_BATCH_SIZE = 128
_NUM_EPOCHS = 250

_TRAINING_ARGS = {
  'dnn_hidden_units': [6, 3],
  'optimizer': tf.keras.optimizers.Adam,
  'optimizer_kwargs': {
      'learning_rate': 0.01
  },
  'layer_activation': None,
  'metrics': ["Accuracy"]
}

_EVAL_METRIC_SPEC = create_metric_spec([
  mse_metric(upper_bound=25, absolute_change=1),
  accuracy_metric()
])

2.特征工程

def preprocessing_fn(inputs):
  """tf.transform's callback function for preprocessing inputs.

  Args:
      inputs: map from feature keys to raw not-yet-transformedfeatures.

  Returns:
      Map from string feature key to transformed feature operations.
  """
  outputs = features_transform(
      inputs=inputs,
      label_key=_LABEL_KEY,
      dense_features=_DENSE_FLOAT_FEATURE_KEYS,
      vocab_features=_VOCAB_FEATURES,
      bucket_features=_BUCKET_FEATURES,
  )
  return outputs

3.建模

def _build_keras_model(**training_args):
   """Build a keras model.

   Args:
       hidden_units: [int], the layer sizes of the DNN (input layer first).
       learning_rate: [float], learning rate of the Adam optimizer.

   Returns:
       A keras model
   """
   feature_columns = \
       prepare_feature_columns(
           dense_features=_DENSE_FLOAT_FEATURE_KEYS,
           vocab_features=_VOCAB_FEATURES,
           bucket_features=_BUCKET_FEATURES,
       )

   return _dnn_regressor(deep_columns=list(feature_columns.values()),
                         dnn_hidden_units=training_args.get(
                             "dnn_hidden_units"),
                         dense_features=_DENSE_FLOAT_FEATURE_KEYS,
                         vocab_features=_VOCAB_FEATURES,
                         bucket_features=_BUCKET_FEATURES,
                         )

4.协调

pipeline_run = WazeTFXPipelineOperator(
   dag=dag,
   task_id='pipeline_run',
   model_name='basic_pipeline_template',
   package=tfx_pipeline_basic,
   pipeline_project_id=EnvConfig.get_value('gcp-project-infra'),
   table_project_id=EnvConfig.get_value('gcp-project-infra'),
   project_utils_filename='utils.py',
   gcp_conn_id=gcp_conn_id,
   enable_pusher=True,
)

很简单，对不对？

配置文件提交到代码库后，系统就会对其进行部署并设置持续训练以及一个完整的流水线，包括所有 TFX 和 Vertex AI 方法，如数据验证、部署到 Dataflow 的转换、监控等。

总结

当我们的一位数据科学家休完长假回来，不得不使用这个新的框架来处理一个案例时，我们知道本项研究取得了不错的结果。据她所说，她能够在几个小时内启动一个可投入生产的完整流水线，而在她休假之前，她需要花几周的时间才能做到这一点。

展望未来，我们计划在 waze-data-tfx 中融入更多内容。我们认为，拥有此通用基础架构的一个关键优势是：添加一个特征后，所有人都可以“免费”使用该特征。例如，我们计划在流水线中添加额外的组件，如 Infra Validator 和 Fairness Indicators。当这些组件得到支持后，每个新的或现有的 ML 流水线将以开箱即用的方式添加这些组件，无需额外的代码。

我们正在计划有关部署的其他改进工作。希望在保证部署质量的同时，尽可能地实现自动化。

我们目前正在探索的一种方式是使用 Canary 部署。数据科学家只需要配置一个评估指标，框架（使用 Vertex Prediction 数据流量分块功能和其他持续评估方法）将在生产环境中测试新模型，并根据评估指标逐步部署或回滚。

点击“阅读原文”访问 TensorFlow 官网

不要忘记“一键三连”哦~

分享

点赞

在看