如何從 ML 庫（例如 PyTorch 或 ZCB20B802A3F0255E0Z251C4E）訪問 GPU 中的 Spark DataFrame 數據

Question

目前我正在研究 Apache Spark 3.0 與 Rapids GPU 加速的使用。 在官方spark-rapids文檔中，我遇到了 這個頁面，其中指出：

在某些情況下，您可能希望訪問 GPU 上的原始數據，最好不要復制它。 一個用例是在進行特征提取后將數據導出到 ML 框架。

對我來說，這聽起來好像可以使 GPU 上已經可用的數據從某些上游 Spark ETL 過程直接可用於 Tensorflow 或 PyTorch 等框架。 如果是這種情況，我如何從這些框架中訪問數據？ 如果我在這里誤解了某些東西，那么這句話到底指的是什么？

Answer 1

您引用的鏈接實際上只允許您訪問仍在 GPU 上的數據，但在另一個框架中使用該數據，如 Tensorflow 或 PyTorch 並不是那么簡單。

TL;博士; 除非您有明確設置的庫以使用 RAPIDS 加速器，否則您可能希望使用 RAPIDS 運行 ETL，然后保存它，並啟動一個新作業以使用該數據訓練您的模型。

您仍然需要解決許多問題。 我們已經在XGBoost的情況下處理了這些問題，但我們還沒有嘗試解決 Tensorflow 或 PyTorch 的問題。

大問題是

1. 獲取數據到正確的過程。 即使數據在 GPU 上，出於安全考慮，它也與給定的用戶進程相關聯。 PyTorch and Tensorflow generally run as python processes and not in the same JVM that Spark is running in. This means that the data has to be sent to the other process. 有幾種方法可以做到這一點，但嘗試將其作為零復制操作來做並非易事。
2. 數據的格式不是 Tensorflow 或 PyTorch 想要的格式。 RAPID 的數據采用箭頭兼容格式。 Tensorflow and PyTorch have APIs for importing data in standard formats from the CPU, but it might take a bit of work to get the data into a format that the frameworks want and to find an API to let you pull it in directly from the GPU.
3. 共享GPU資源。 Spark 最近才添加了對調度 GPU 的支持。 在此之前，人們只需為每個執行程序啟動一個 spark 任務和一個 python 進程，以便 python 進程在進行訓練或推理時擁有整個 GPU。 借助 RAPIDS 加速器，GPU 不再免費，您需要一種共享資源的方法。 如果兩個庫都更新為使用它並且它們處於同一進程中，則 RMM 會提供其中的一些，但在 Pytorch 和 Tensoflow 的情況下，它們通常在 python 進程中，因此弄清楚如何共享 Z52F9EC21735243AD9917CDA3CA077 是困難的。