機(jī)器學(xué)習(xí)中，有一個(gè)限制速度的環(huán)節(jié)，那就是從 CPU 到 GPU 之間的張量遷移。很多計(jì)算只能在 CPU 上進(jìn)行，然后遷移到 GPU 進(jìn)行后續(xù)的訓(xùn)練工作，因此遷移中如果速度太慢，則會(huì)拖累整個(gè)模型的訓(xùn)練效率。近日，有一位開發(fā)者開源了針對(duì) PyTorch 的 CPU->GPU 遷移工具，相比原版加速了 110 倍之多。
神經(jīng)網(wǎng)絡(luò)的訓(xùn)練中往往需要進(jìn)行很多環(huán)節(jié)的加速，這就是為什么我們逐漸使用 GPU 替代 CPU、使用各種各樣的算法來加速機(jī)器學(xué)習(xí)過程。但是，在很多情況下，GPU 并不能完成 CPU 進(jìn)行的很多操作。比如訓(xùn)練詞嵌入時(shí)，計(jì)算是在 CPU 上進(jìn)行的，然后需要將訓(xùn)練好的詞嵌入轉(zhuǎn)移到 GPU 上進(jìn)行訓(xùn)練。
在這一過程中，張量遷移可能會(huì)變得非常緩慢，這成為了機(jī)器學(xué)習(xí)訓(xùn)練的一個(gè)瓶頸。
面對(duì)這樣的問題，在很多優(yōu)化方法中，就是將操作盡可能放在 GPU 上（如直接在 GPU 上進(jìn)行數(shù)據(jù)預(yù)處理、詞嵌入等的操作），并努力減少兩者之間的數(shù)據(jù)交互，因?yàn)檫@些環(huán)節(jié)都很費(fèi)時(shí)。機(jī)器之心就曾報(bào)道過這樣的一個(gè)教程，將數(shù)據(jù)預(yù)處理放在了 GPU 上，減少了不少時(shí)間。
以上事例說明，如果能夠做好 CPU 和 GPU 之間的遷移，則可以幫助開發(fā)者更好地優(yōu)化機(jī)器學(xué)習(xí)模型，使 CPU、GPU 等硬件更好地完成自己的工作。
近日，有一位開發(fā)者就開源了一個(gè)名為 SpeedTorch 的工具。這一工具庫可以實(shí)現(xiàn)高達(dá) 110 倍的 CPU 到 GPU 遷移加速。
項(xiàng)目地址：https://github.com/Santosh-Gupta/SpeedTorch
項(xiàng)目背景
作者表示，最初想要?jiǎng)?chuàng)建 SpeedTorch 庫是為了幫助訓(xùn)練大量的嵌入向量，而 GPU 在 RAM 保存這些嵌入時(shí)可能有困難。為了解決這個(gè)問題，他發(fā)現(xiàn)在 CPU 上托管一些此類嵌入有助于在 GPU 上保存它們。嵌入系統(tǒng)采用稀疏訓(xùn)練，只有一部分參數(shù)參與前饋/更新操作，剩余參數(shù)處于閑置狀態(tài)。所以作者想到，為什么不在訓(xùn)練的過程中關(guān)閉這些閑置參數(shù)呢？這就需要快速的 CPU→GPU 數(shù)據(jù)遷移能力。
隨著 CPU→GPU 遷移速度的加快，除了加速了 CPU 到 GPU 的張量轉(zhuǎn)移外，開發(fā)者還可以實(shí)現(xiàn)很多新的功能。
- 將 SpeedTorch 庫嵌入數(shù)據(jù)管道中，實(shí)現(xiàn) CPU 和 GPU 之間快速的雙向數(shù)據(jù)交互；
- 通過 CPU 存儲(chǔ)將模型的訓(xùn)練參數(shù)增加近兩倍（閑置參數(shù)保存在 CPU 中，需要更新時(shí)再移動(dòng)到 GPU 里，因此可以擴(kuò)大模型整體的參數(shù)量）；
- 在訓(xùn)練稀疏嵌入向量中采用 Adadelta、Adamax、RMSprop、Rprop、ASGD、AdamW 和 Adam 優(yōu)化器。之前只有 SpraseAdam、Adagrad 和 SGD 適合稀疏嵌入訓(xùn)練。
那么，能夠?qū)崿F(xiàn)如此驚人的加速的庫是怎么實(shí)現(xiàn)的呢？
SpeedTorch
背后的技術(shù)
SpeedTorch 如此之快的技術(shù)是因?yàn)樗腔?Cupy 開發(fā)的。CuPy 是一個(gè)借助 CUDA GPU 庫在英偉達(dá) GPU 上實(shí)現(xiàn) Numpy 數(shù)組的庫?；?Numpy 數(shù)組的實(shí)現(xiàn)，GPU 自身具有的多個(gè) CUDA 核心可以促成更好的并行加速。
CuPy 接口是 Numpy 的一個(gè)鏡像，并且在大多情況下，它可以直接替換 Numpy 使用。只要用兼容的 CuPy 代碼替換 Numpy 代碼，用戶就可以實(shí)現(xiàn) GPU 加速。
CuPy 支持 Numpy 的大多數(shù)數(shù)組運(yùn)算，包括索引、廣播、數(shù)組數(shù)學(xué)以及各種矩陣變換。
有了這樣強(qiáng)大的底層支持，再加上一些優(yōu)化方法，SpeedTorch 就能達(dá)到 110 倍的速度了。
使用方法
SpeedTorch 可以通過 pip 安裝。你需要在導(dǎo)入 SpeedTorch 之前事先安裝和導(dǎo)入 Cupy。
安裝步驟如下：
利用 SpeedTorch 加快 CPU→GPU 數(shù)據(jù)遷移速度
如下 colab notebook 所示，如何利用 Data Gadget 將數(shù)據(jù)載入 SpeedTorch，以及如何將數(shù)據(jù)移入/移出 Pytorch cuda 變量。
代碼示例：https://colab.research.google.com/drive/185Z5Gi62AZxh-EeMfrTtjqxEifHOBXxF
借助于 SpeedTorch 將非稀疏優(yōu)化器（本例中為 Adamax）用于稀疏訓(xùn)練
代碼示例: https://colab.research.google.com/drive/1ApJR3onbgQWM3FBcBKMvwaGXIDXlDXOt
以上展示了如何以常規(guī)的方式訓(xùn)練 word2vec，隨后展示了如何使用 SpeedTorch 在同樣的數(shù)據(jù)上進(jìn)行訓(xùn)練——在通常不支持稀疏訓(xùn)練的優(yōu)化器上。因?yàn)榍度胱兞堪乃星度朐诿恳徊可隙加懈?，你可以在初始化期間將 sparse=False。
效果
這一部分記錄了 Cupy/PyTorch 張量和 PyTorch 變量之間的數(shù)據(jù)遷移速度。其中，需要遷移 128 維的嵌入向量，共有 131,072 個(gè) 32 位浮點(diǎn)數(shù)。使用了如下的代碼進(jìn)行測(cè)試工作。所有測(cè)試都使用了特斯拉 K80 GPU。
測(cè)試代碼鏈接：https://colab.research.google.com/drive/1b3QpfSETePo-J2TjyO6D2LgTCjVrT1lu
下表是結(jié)果摘要。在同樣情況下，將數(shù)據(jù)從 PyTorch CUDA 張量傳遞到 CUDA PyTorch 嵌入變量上是要比 SpeedTorch 更快的，但對(duì)于所有其他的傳輸類型，SpeedTorch 更快。對(duì)于轉(zhuǎn)移到 Cuda Pytorch 嵌入，或從 Cuda Pytorch 嵌入轉(zhuǎn)移的兩個(gè)步驟的總和上來說，SpeedTorch 比常規(guī) GPU 和 CPU Pinned 張量的 Pytorch 速度同樣快。
從表中可以看出，這是 SpeedTorch 確實(shí)比 PyTorch 自帶的數(shù)據(jù)遷移方法要快很多。
【責(zé)任編輯：張燕妮 TEL：（010）68476606】