Spark-TFRecord：Apache Spark与TensorFlow TFRecord互操作示例

Posted 2021-04-20 不靠谱的猫

在真实的机器学习场景中，TensorFlow和Apache Spark的互操作是一个常见的挑战。TensorFlow可以说是市场上最受欢迎的深度学习框架，而Apache Spark仍然是使用最广泛的数据计算平台之一。尽管存在使TensorFlow适应Spark的框架，但是互操作性通常是在数据级别。TFRecord是TensorFlow中的native数据结构，Apache Spark不完全支持它。最近，LinkedIn的工程师们公开了Spark-TFRecord，这是Spark基于TensorFlow TFRecord的一个新的本地数据源。

LinkedIn决定解决这个问题并不奇怪。这家互联网巨头长期以来一直广泛采用Spark技术，并一直是TensorFlow和机器学习开源社区的积极贡献者。在内部，LinkedIn的工程团队经常尝试在TensorFlow的TFRecord格式和Spark的Avro或Parquet等格式之间进行转换。Spark-TFRecord项目的目标是在Spark管道中提供TFRecord结构的native功能。

之前的方法

Spark- tfrecord并不是第一个尝试解决Spark和TensorFlow之间的数据互操作性挑战的项目。其中最受欢迎的项目是Spark的创建者Databricks推广的Spark- tensorflow - connector。在体系结构上，connector是将TFRecord格式适应为Spark SQL数据流。Spark-TensorFlow-Connector在关系数据访问场景中非常有效，但在其他用例中效果就非常有限。

TensorFlow工作流的一个重要部分是与磁盘I/O操作相关的，而不是与数据库访问相关的。在这些场景中，当使用Spark-TensorFlow-Connector时，开发人员最终要编写大量的代码。此外，Spark-TensorFlow-Connector的当前版本仍然缺乏重要的功能，比如在TensorFlow计算中经常使用的Partition。另外，connector更像是在Spark SQL Data Frames中处理TensorFlow records的桥梁，而不是本地文件格式。

为了克服这些限制，LinkedIn工程团队决定从不同的角度来解决Spark-TensorFlow互操作性问题。

Spark-TFRecord

Spark-TFRecord是Apache Spark的本地TensorFlow TFRecord。具体来说，Spark-TFRecord提供了用于Apache Spark和TFREcord数据互操作的例程。与构建connector来处理TFRecord结构不同，Spark-TFRecord被构建为一个本地Spark数据集，就像Avro、JSON或Parquet一样。这意味着所有Spark的DataSet和DataFrame I / O例程都可以在Spark-TFRecord中自动使用。

LinkedIn的工程团队实现一个新的Spark FileFormat接口，该接口的基本设计目的是支持磁盘I/O操作。新的接口将使TFRecord本地操作适应于任何Spark数据流。在架构上，Spark-TFRecord由一系列基本的构建模块组成:

• Schema Inferencer：这是最接近Spark-TensorFlow-Connector的组件。此接口将TFRecords表示映射到本机Spark数据类型。

• TFRecord Reader：该组件读取TFRecord结构并将其传递给反序列化器（deserializer）。

• TFRecord Writer：此组件从序列化器（serializer）接收TFRecord结构并将其写入磁盘。

• TFRecord Serializer：此组件将Spark InternalRow转换为TFRecord结构。

• TFRecord Deserializer：此组件将TFRecords转换为Spark InternalRow结构。

使用LinkedIn的Spark- tfrecord与其他Spark native数据集没有什么不同。开发人员只需包含spark-tfrecord jar库，并使用传统的DataFrame API来读写TFRecords，如下面的代码所示:

对于大多数组织而言，Spark与TensorFlow等深度学习框架之间的互操作性可能仍将是一个充满挑战的领域。但是，像LinkedIn的Spark-TFRecord这样经过大量测试的项目无疑有助于简化这两种技术之间的互操作性，而这两种技术对于许多现代机器学习架构来说都是必不可少的。