如何从 ML 库(如 PyTorch 或 Tensorflow)访问 GPU 中的 Spark DataFrame 数据

How to access Spark DataFrame data in GPU from ML Libraries such as PyTorch or Tensorflow

目前我正在研究 Apache Spark 3.0 与 Rapids GPU 加速的使用。在官方 spark-rapids 文档中,我遇到了 this page 其中指出:

There are cases where you may want to get access to the raw data on the GPU, preferably without copying it. One use case for this is exporting the data to an ML framework after doing feature extraction.

对我来说,这听起来好像可以使来自某些上游 Spark ETL 过程的 GPU 上已经可用的数据直接可用于 Tensorflow 或 PyTorch 等框架。如果是这种情况,我如何从这些框架中访问数据?如果我在这里误解了什么,引述到底指的是什么?

您引用的 link 实际上只允许您访问仍位于 GPU 上的数据,但在其他框架(如 Tensorflow 或 PyTorch)中使用该数据并不是那么简单。

TL;DR;除非你明确设置了一个库来使用 RAPIDS 加速器,否则你可能想 运行 使用 RAPIDS 你的 ETL,然后保存它,并启动一个新的工作来训练您的模型使用该数据。

您还有许多问题需要解决。我们已经在 XGBoost 的情况下处理了这些问题,但我们还没有尝试为 Tensorflow 或 PyTorch 解决这些问题。

主要问题是

  1. 将数据传送到正确的进程。即使数据在 GPU 上,出于安全考虑,它也会绑定到给定的用户进程。 PyTorch 和 Tensorflow 通常 运行 作为 python 进程,而不是在与 Spark 运行 相同的 JVM 中。这意味着数据必须发送到其他进程。有几种方法可以做到这一点,但尝试将其作为零复制操作进行操作并非易事。
  2. 数据的格式不是Tensorflow或PyTorch想要的。 RAPIDs 的数据采用箭头兼容格式。 Tensorflow 和 PyTorch 有 APIs 用于从 CPU 以标准格式导入数据,但可能需要一些工作才能将数据转换为框架所需的格式并找到 API 让您直接从 GPU 中提取它。
  3. 共享 GPU 资源。 Spark 最近才添加了对调度 GPU 的支持。在此之前,人们只会为每个执行程序启动一个 spark 任务和一个 python 进程,这样 python 进程在进行训练或推理时将拥有整个 GPU。使用 RAPIDS 加速器,GPU 不再是免费的,您需要一种共享资源的方法。如果两个库都更新为使用它并且它们在同一进程中,RMM 会提供其中的一些功能,但是对于 Pytorch 和 Tensoflow,它们通常在 python 个进程中,因此很难弄清楚如何共享 GPU .