【问题标题】：How to access Spark DataFrame data in GPU from ML Libraries such as PyTorch or Tensorflow如何从 ML 库（如 PyTorch 或 Tensorflow）访问 GPU 中的 Spark DataFrame 数据
【发布时间】：2021-01-04 16:02:29
【问题描述】：
目前我正在研究使用 Rapids GPU 加速的 Apache Spark 3.0。在官方spark-rapids 文档中，我遇到了this page，其中指出：
在某些情况下，您可能希望访问 GPU 上的原始数据，最好不要复制它。一个用例是在进行特征提取后将数据导出到 ML 框架。
在我看来，这听起来好像可以让来自某些上游 Spark ETL 进程的 GPU 上已经可用的数据直接可用于诸如 Tensorflow 或 PyTorch 之类的框架。如果是这种情况，我如何从这些框架中访问数据？如果我在这里误解了某些东西，那么引用的确切含义是什么？
【问题讨论】：
标签：
tensorflow
apache-spark
pytorch
rapids
spark3
【解决方案1】：
您引用的链接实际上只允许您访问仍在 GPU 上的数据，但在另一个框架（如 Tensorflow 或 PyTorch）中使用该数据并不是那么简单。
TL;DR;除非您有明确设置的库以使用 RAPIDS 加速器，否则您可能希望使用 RAPIDS 运行 ETL，然后保存它，并启动一个新作业以使用该数据训练您的模型。
仍有许多问题需要解决。我们已经在 XGBoost 的情况下解决了这些问题，但我们还没有尝试为 Tensorflow 或 PyTorch 解决这些问题。
最大的问题是
- 将数据传送到正确的进程。即使数据在 GPU 上，出于安全考虑，它也与给定的用户进程相关联。 PyTorch 和 Tensorflow 通常作为 python 进程运行，而不是在运行 Spark 的同一个 JVM 中。这意味着必须将数据发送到另一个进程。有几种方法可以做到这一点，但尝试将其作为零复制操作来做并非易事。
- 数据的格式不是 Tensorflow 或 PyTorch 想要的。 RAPID 的数据采用箭头兼容格式。 Tensorflow 和 PyTorch 具有用于从 CPU 以标准格式导入数据的 API，但可能需要一些工作才能将数据转换为框架所需的格式，并找到一个 API 让您直接从 GPU 中提取数据。
- 共享 GPU 资源。 Spark 最近才添加了对调度 GPU 的支持。在此之前，人们只需为每个执行程序和一个 python 进程启动一个 spark 任务，以便在进行训练或推理时，python 进程将拥有整个 GPU。使用 RAPIDS 加速器，GPU 不再是免费的，您需要一种共享资源的方法。如果两个库都更新为使用 RMM 并且它们在同一个进程中，则 RMM 会提供其中的一部分，但在 Pytorch 和 Tensoflow 的情况下，它们通常在 python 进程中，因此很难弄清楚如何共享 GPU。