Share via

使用 Petastorm 加载数据

  • 项目

本文介绍如何使用 Petastorm 将数据从 Apache Spark 转换为 TensorFlow 或 PyTorch。 它还提供了一个示例，可演示如何使用 Petastorm 为 ML 准备数据。

Petastorm 是一种开放源代码数据访问库。 它可直接从采用 Apache Parquet 格式的数据集和已加载为 Apache Spark 数据帧的数据集对深度学习模型启用单节点或分布式训练和评估。 Petastorm 支持基于 Python 的热门机器学习 (ML) 框架，如 Tensorflow、PyTorch，和 PySpark。 有关 Petastorm 的详细信息，请参阅 Petastorm API 文档。

使用 Petastorm 从 Spark 数据帧加载数据

Petastorm Spark 转换器 API 简化了从 Spark 到 TensorFlow 或 PyTorch 的数据转换。 输入 Spark 数据帧首先以 Parquet 格式具体化，然后作为 tf.data.Dataset 或 torch.utils.data.DataLoader 进行加载。 请参阅 Petastorm API 文档中的 Spark 数据集转换器 API 部分。

建议的工作流为：

  1. 使用 Apache Spark 来加载数据，还可以选择对数据进行预处理。
  2. 使用 Petastorm spark_dataset_converter 方法可将 Spark 数据帧中的数据转换为 TensorFlow 数据集或 PyTorch DataLoader。
  3. 将数据传入 DL 框架进行训练或推理。

配置缓存目录

Petastorm Spark 转换器以 Parquet 格式将输入 Spark 数据帧缓存到用户指定的缓存目录位置。 缓存目录必须是以 file:///dbfs/ 开头的 DBFS 路径，例如 file:///dbfs/tmp/foo/ 表示 dbfs:/tmp/foo/ 的相同位置。 可以通过两种方式配置缓存目录：

  • 在群集 Spark 配置中添加行：petastorm.spark.converter.parentCacheDirUrl file:///dbfs/...

  • 在笔记本中，调用 spark.conf.set()：

    from petastorm.spark import SparkDatasetConverter, make_spark_converter

spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'file:///dbfs/...')

在使用缓存后，可以通过调用 converter.delete() 来显式删除缓存，或通过在对象存储中配置生命周期规则隐式管理缓存。

Databricks 支持三种方案中的 DL 训练：

  • 单节点训练
  • 分布式超参数优化
  • 分布式训练

有关端到端示例，请参阅以下笔记本：

使用 Petastorm 直接加载 Parquet 文件

此方法优先级低于 Petastorm Spark 转换器 API。

建议的工作流为：

  1. 使用 Apache Spark 来加载数据，还可以选择对数据进行预处理。
  2. 将 Parquet 格式的数据保存到具有随附 DBFS 装载的 DBFS 路径。
  3. 通过 DBFS 装入点以 Petastorm 格式加载数据。
  4. 在 DL 框架中使用数据进行训练或推理。

有关端到端示例，请参阅示例笔记本。

示例：使用 TensorFlow 或 PyTorch 预处理数据并训练模型

此示例笔记本演示了 Databricks 上的以下工作流：

  1. 使用 Spark 加载数据。
  2. 使用 Petastorm 将 Spark 数据帧转换为 TensorFlow 数据集。
  3. 将数据馈送到单节点 TensorFlow 模型中以进行训练。
  4. 将数据馈送到分布式超参数优化函数中。
  5. 将数据馈送到分布式 TensorFlow 模型中进行训练。

简化从 Spark 到 TensorFlow 笔记本的数据转换

获取笔记本

此示例笔记本演示了 Databricks 上的以下工作流：

  1. 使用 Spark 加载数据。
  2. 使用 Petastorm 将 Spark 数据帧转换为 PyTorch DataLoader。
  3. 将数据馈送到单节点 PyTorch 模型中进行训练。
  4. 将数据馈送到分布式超参数优化函数中。
  5. 将数据馈送到分布式 PyTorch 模型中进行训练。

简化从 Spark 到 PyTorch 笔记本的数据转换

获取笔记本

示例：使用 Petastorm 预处理数据并加载 Parquet 文件

此示例笔记本显示了 Databricks 上的以下工作流：

  1. 使用 Spark 加载和预处理数据。
  2. 在 dbfs:/ml 使用 Parquet 保存数据。
  3. 通过优化的 FUSE 装载 file:/dbfs/ml 使用 Petastorm 加载数据。
  4. 将数据馈送到深度学习框架中进行训练或推理。

使用 Spark 和 Petastorm 为深度学习笔记本准备数据

获取笔记本