使用RAPIDS在NVIDIA GPU上分析腦細胞基準

單細胞基因組學(xué)研究繼續(xù)推進疾病預(yù)防藥物的發(fā)現(xiàn)。例如，它在為當前的新冠肺炎大流行開發(fā)治療、識別易受感染的細胞以及揭示受感染患者免疫系統(tǒng)的變化方面起著關(guān)鍵作用。然而，隨著大規(guī)模單細胞數(shù)據(jù)集可用性的不斷提高，計算效率的低下明顯影響了科學(xué)研究的速度。將這些計算瓶頸轉(zhuǎn)移到 GPU 已經(jīng)證明了有趣的結(jié)果。

在最近的一篇博客文章中， NVIDIA 對 100 萬個小鼠腦細胞進行了基準分析，這些腦細胞通過 10 倍基因組學(xué)測序。結(jié)果表明，在 GCP CPU 實例上運行端到端工作流需要三個多小時，而在單個 NVIDIA V100 GPU 上處理整個數(shù)據(jù)集只需 11 分鐘。此外，在 GCP GPU 實例上運行 RAPIDS 分析的成本也比 CPU 版本低 3 倍。此處閱讀博客。

按照Jupyter 筆記本對該數(shù)據(jù)集進行 RAPIDS 分析。要運行筆記本，文件rapids_scanpy_funcs.py和utils.py必須與筆記本位于同一文件夾中。我們提供了第二個筆記本，其中包含 CPU 版本的分析here。在與 Google Dataproc 團隊的合作下，我們構(gòu)建了一個入門指南，以幫助開發(fā)人員快速運行這個轉(zhuǎn)錄組學(xué)用例。最后，看看這個 NVIDIA 和谷歌云共同撰寫博客文章，它展示了工作的影響。

對 GPU 進行單細胞 RNA 分析

執(zhí)行單細胞分析的典型工作流程通常從一個矩陣開始，該矩陣映射每個細胞中測量的每個基因腳本的計數(shù)。執(zhí)行預(yù)處理步驟以濾除噪聲，并對數(shù)據(jù)進行歸一化以獲得在數(shù)據(jù)集的每個單獨單元中測量的每個基因的表達。在這一步中，機器學(xué)習(xí)也常用于糾正數(shù)據(jù)收集中不需要的偽影?；虻臄?shù)量通常相當大，這會產(chǎn)生許多不同的變異，并在計算細胞之間的相似性時增加很多噪音。在識別和可視化具有相似基因表達的細胞簇之前，特征選擇和降維可以減少這種噪聲。這些細胞簇的轉(zhuǎn)錄表達也可以進行比較，以了解為什么不同類型的細胞行為和反應(yīng)不同。

圖 1 ：顯示單細胞 RNA 測序數(shù)據(jù)分析步驟的管道。從每個細胞中的基因活性矩陣開始， RAPIDS 文庫可用于將矩陣轉(zhuǎn)換為基因表達，對細胞進行聚類和布局以供可視化，并幫助分析具有不同活性的基因。

該分析證明了使用RAPIDS加速使用單個 GPU 分析 100 萬個細胞的單細胞 RNA 序列數(shù)據(jù)。然而，實驗只處理了前 100 萬個細胞，而不是整個 130 萬個細胞。因此，在單細胞 RNA 數(shù)據(jù)的工作流中處理所有 130 萬個細胞的時間幾乎是單個 V100 GPU 的兩倍。另一方面，相同的工作流在單個 NVIDIA A100 40GB GPU 上只需 11 分鐘。不幸的是， V100 的性能下降了近 2 倍，主要原因是 GPU 的內(nèi)存被超額訂閱，從而在需要時溢出到主機內(nèi)存。在下一節(jié)中，我們將更詳細地介紹這種行為，但需要明確的是， GPU 的內(nèi)存是擴展的限制因素。因此，更快地處理更大的工作負載需要更強大的 GPU 服務(wù)器，如 A100 或/或?qū)⑻幚矸稚⒌蕉鄠€ GPU 服務(wù)器上。

將預(yù)處理擴展到多個 GPU 的好處

當工作流的內(nèi)存使用量超過單個 GPU 的容量時，統(tǒng)一虛擬內(nèi)存（ UVM ）可用于超額訂閱 GPU ，并自動溢出到主內(nèi)存。這種方法在探索性數(shù)據(jù)分析過程中是有利的，因為適度的超額訂閱率可以消除在 GPU 內(nèi)存不足時重新運行工作流的需要。

但是，嚴格依靠 UVM 將 GPU 的內(nèi)存超額訂閱 2 倍或更多可能會導(dǎo)致性能不佳。更糟糕的是，當任何單個計算需要的內(nèi)存超過 NVIDIA 上的可用內(nèi)存時，它可能會導(dǎo)致執(zhí)行無限期掛起。將計算擴展到多個 GPU 可以提高并行性并減少每個 GPU 上的內(nèi)存占用。在某些情況下，它可以消除超額認購的需要。圖 2 表明，我們可以通過將預(yù)處理計算擴展到多個 GPU 來實現(xiàn)線性縮放，與單個 GPU V100 GPU 相比， 8 個 GPU s 會產(chǎn)生略微超過 8 倍的加速比?？紤]到這一點，需要不到 2 分鐘才能將 130 萬個細胞和 18k 基因的數(shù)據(jù)集減少到約 129 萬個細胞和 8 GPU上 4k 個高度可變的基因。這超過了 8 。 55 倍的加速，因為單個 V100 需要 16 分鐘來運行相同的預(yù)處理步驟。

圖 2 ：具有不同硬件配置的 130 萬小鼠腦細胞上典型單細胞 RNA 工作流的運行時間（秒）比較。在 GPU 上執(zhí)行這些計算表明性能大幅提高。

圖 3 ：單個 GPU 配置的運行時主要由預(yù)處理步驟控制，在單個 V100 上占據(jù) 75% 的端到端運行時，在單個 A100 上占據(jù) 70% 的運行時。利用 DGX1 上的所有 GPU 將比率降低到略高于 32%.

使用 Dask 和 RAPIDS 將單細胞 RNA 筆記本擴展到多個 GPU

許多預(yù)處理步驟，如加載數(shù)據(jù)集、過濾嘈雜的轉(zhuǎn)錄本和細胞、將計數(shù)標準化為表達式以及特征選擇，本質(zhì)上是并行的，每個 GPU 獨立負責(zé)其子集。糾正數(shù)據(jù)收集噪音影響的一個常見步驟是使用不需要的基因（如核糖體基因）的貢獻比例，并擬合許多小型線性回歸模型，數(shù)據(jù)集中每個轉(zhuǎn)錄本對應(yīng)一個模型。由于轉(zhuǎn)錄本的數(shù)量通?？梢赃_到數(shù)萬個，因此通常使用分散性或可變性的度量方法，只選擇幾千個最具代表性的基因。

Dask是一個優(yōu)秀的庫，用于在一組工作進程上分發(fā)數(shù)據(jù)處理工作流。 RAPIDS 通過將每個工作進程映射到自己的 GPU ，使 Dask 也能夠使用 GPU s 。此外， Dask 提供了一個分布式陣列對象，非常類似于 NumPy 陣列的分布式版本（或CuPy，其 GPU 加速外觀相似），它允許用戶在多個 GPU 上，甚至跨多臺物理機器，分發(fā)上述預(yù)處理操作的步驟，操作和轉(zhuǎn)換數(shù)據(jù)的方式與 NumPy 或 CuPy 數(shù)組大致相同。

在預(yù)處理之后，我們還通過對數(shù)據(jù)子集進行訓(xùn)練并分配推理來分配主成分分析（ PCA ）縮減步驟，通過僅將前 50 個主成分恢復(fù)到單個 GPU 來降低通信成本，用于剩余的聚類和可視化步驟。該數(shù)據(jù)集的 PCA 簡化單元矩陣僅為 260 MB ，允許在單個 GPU 上執(zhí)行剩余的聚類和可視化步驟。使用這種設(shè)計，即使包含 500 萬個單元的數(shù)據(jù)集也只需要 1GB 內(nèi)存。

圖 4 ： 1 。 3M 小鼠腦細胞的樣本可視化，使用 cuML 的 UMAP 縮小為二維，并使用 cuGraph 的 Louvain 聚集

結(jié)論

以我們計算工具的發(fā)展速度，我們可以假設(shè)數(shù)據(jù)處理量很快就會迎頭趕上，特別是對于單細胞分析工作負載，這迫使我們需要更高的擴展。同時，通過將聚類和可視化步驟分布在多個 RAPIDS 上，仍有機會進一步減少探索性數(shù)據(jù)分析過程的迭代次數(shù)。更快的迭代意味著更好的模型，縮短洞察時間，以及更明智的結(jié)果。除 T-SNE 外，多 GPU 工作流筆記本的所有集群和可視化步驟都可以通過 GPU cuML 和 cuGraph 分布在 GPU 上的 Dask 工作人員上。

關(guān)于作者

Corey Nolet 是 NVIDIA 的 RAPIDS ML 團隊的數(shù)據(jù)科學(xué)家兼高級工程師，他專注于構(gòu)建和擴展機器學(xué)習(xí)算法，以支持光速下的極端數(shù)據(jù)負載。在 NVIDIA 工作之前， Corey 花了十多年時間為國防工業(yè)的 HPC 環(huán)境構(gòu)建大規(guī)模探索性數(shù)據(jù)科學(xué)和實時分析平臺?？评锍钟杏砉W(xué)士學(xué)位計算機科學(xué)碩士學(xué)位。他還在攻讀博士學(xué)位。在同一學(xué)科中，主要研究圖形和機器學(xué)習(xí)交叉點的算法加速。科里熱衷于利用數(shù)據(jù)更好地了解世界。

審核編輯：郭婷