使用RAPIDS在NVIDIA GPU上分析腦細(xì)胞基準(zhǔn)

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

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

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

對(duì) GPU 進(jìn)行單細(xì)胞 RNA 分析

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

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

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

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

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

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

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

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

使用 Dask 和 RAPIDS 將單細(xì)胞 RNA 筆記本擴(kuò)展到多個(gè) GPU

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

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

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

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

結(jié)論

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

關(guān)于作者

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

審核編輯：郭婷