選自Github
作者:Santosh Gupta
機器之心編譯
參與:杜偉、一鳴、澤南
機器學(xué)習(xí)中,有一個限制速度的環(huán)節(jié),那就是從 CPU 到 GPU 之間的張量遷移。很多計算只能在 CPU 上進行,然后遷移到 GPU 進行后續(xù)的訓(xùn)練工作,因此遷移中如果速度太慢,則會拖累整個模型的訓(xùn)練效率。近日,有一位開發(fā)者開源了針對 PyTorch 的 CPU->GPU 遷移工具,相比原版加速了 110 倍之多。
神經(jīng)網(wǎng)絡(luò)的訓(xùn)練中往往需要進行很多環(huán)節(jié)的加速,這就是為什么我們逐漸使用 GPU 替代 CPU、使用各種各樣的算法來加速機器學(xué)習(xí)過程。但是,在很多情況下,GPU 并不能完成 CPU 進行的很多操作。比如訓(xùn)練詞嵌入時,計算是在 CPU 上進行的,然后需要將訓(xùn)練好的詞嵌入轉(zhuǎn)移到 GPU 上進行訓(xùn)練。
在這一過程中,張量遷移可能會變得非常緩慢,這成為了機器學(xué)習(xí)訓(xùn)練的一個瓶頸。
面對這樣的問題,在很多優(yōu)化方法中,就是將操作盡可能放在 GPU 上(如直接在 GPU 上進行數(shù)據(jù)預(yù)處理、詞嵌入等的操作),并努力減少兩者之間的數(shù)據(jù)交互,因為這些環(huán)節(jié)都很費時。機器之心就曾報道過這樣的一個教程,將數(shù)據(jù)預(yù)處理放在了 GPU 上,減少了不少時間。
以上事例說明,如果能夠做好 CPU 和 GPU 之間的遷移,則可以幫助開發(fā)者更好地優(yōu)化機器學(xué)習(xí)模型,使 CPU、GPU 等硬件更好地完成自己的工作。
近日,有一位開發(fā)者就開源了一個名為 SpeedTorch 的工具。這一工具庫可以實現(xiàn)高達(dá) 110 倍的 CPU 到 GPU 遷移加速。
項目地址:https://github.com/Santosh-Gupta/SpeedTorch
項目背景
作者表示,最初想要創(chuàng)建 SpeedTorch 庫是為了幫助訓(xùn)練大量的嵌入向量,而 GPU 在 RAM 保存這些嵌入時可能有困難。為了解決這個問題,他發(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ā)者還可以實現(xiàn)很多新的功能。
將 SpeedTorch 庫嵌入數(shù)據(jù)管道中,實現(xiàn) CPU 和 GPU 之間快速的雙向數(shù)據(jù)交互;
通過 CPU 存儲將模型的訓(xùn)練參數(shù)增加近兩倍(閑置參數(shù)保存在 CPU 中,需要更新時再移動到 GPU 里,因此可以擴大模型整體的參數(shù)量);
在訓(xùn)練稀疏嵌入向量中采用 Adadelta、Adamax、RMSprop、Rprop、ASGD、AdamW 和 Adam 優(yōu)化器。之前只有 SpraseAdam、Adagrad 和 SGD 適合稀疏嵌入訓(xùn)練。
那么,能夠?qū)崿F(xiàn)如此驚人的加速的庫是怎么實現(xiàn)的呢?
SpeedTorch
背后的技術(shù)
SpeedTorch 如此之快的技術(shù)是因為它是基于 Cupy 開發(fā)的。CuPy 是一個借助 CUDA GPU 庫在英偉達(dá) GPU 上實現(xiàn) Numpy 數(shù)組的庫?;?Numpy 數(shù)組的實現(xiàn),GPU 自身具有的多個 CUDA 核心可以促成更好的并行加速。
CuPy 接口是 Numpy 的一個鏡像,并且在大多情況下,它可以直接替換 Numpy 使用。只要用兼容的 CuPy 代碼替換 Numpy 代碼,用戶就可以實現(xiàn) GPU 加速。
CuPy 支持 Numpy 的大多數(shù)數(shù)組運算,包括索引、廣播、數(shù)組數(shù)學(xué)以及各種矩陣變換。
有了這樣強大的底層支持,再加上一些優(yōu)化方法,SpeedTorch 就能達(dá)到 110 倍的速度了。
使用方法
SpeedTorch 可以通過 pip 安裝。你需要在導(dǎo)入 SpeedTorch 之前事先安裝和導(dǎo)入 Cupy。
安裝步驟如下:
!pip install SpeedTorch
import cupy
import SpeedTorch
利用 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)練
SkipGram_ModelRegular = SkipGramModelRegular(numEmbeds=number_items, emb_dimension=128, sparseB=True)
use_cuda = torch.cuda.is_available()
if use_cuda:
SkipGram_ModelRegular.cuda()
optimizer = optim.SparseAdam(
SkipGram_ModelRegular.parameters())
runningLoss = 0
runnngTime = 0
batch_size = 512
negSamp = 64
numPos = 4
skip_window = int(numPos/2)
targets = torch.ones( batch_size, numPos + negSamp , dtype = torch.float32 ).cuda()
for i in range(500):
batch, labels, negz = generate_batch(batch_size=batch_size, skip_window=skip_window, negRate= negSamp)
batchTensor = torch.from_numpy(batch)
LabelTensor = torch.from_numpy(labels)
negTensor = torch.from_numpy(negz)
pos_u = Variable(torch.LongTensor(LabelTensor.long()))
pos_v = Variable(torch.LongTensor(batchTensor.long()))
neg_v = Variable(torch.LongTensor(negTensor.long()))
if use_cuda:
pos_u = pos_u.cuda()
pos_v = pos_v.cuda()
neg_v = neg_v.cuda()
optimizer.zero_grad()
loss = SkipGram_ModelRegular.forward(pos_u, pos_v, neg_v, targets)
runningLoss = runningLoss + loss.data.item()
loss.backward()
optimizer.step()
代碼示例: https://colab.research.google.com/drive/1ApJR3onbgQWM3FBcBKMvwaGXIDXlDXOt
以上展示了如何以常規(guī)的方式訓(xùn)練 word2vec,隨后展示了如何使用 SpeedTorch 在同樣的數(shù)據(jù)上進行訓(xùn)練——在通常不支持稀疏訓(xùn)練的優(yōu)化器上。因為嵌入變量包含的所有嵌入在每一部上都有更新,你可以在初始化期間將 sparse=False。
效果
這一部分記錄了 Cupy/PyTorch 張量和 PyTorch 變量之間的數(shù)據(jù)遷移速度。其中,需要遷移 128 維的嵌入向量,共有 131,072 個 32 位浮點數(shù)。使用了如下的代碼進行測試工作。所有測試都使用了特斯拉 K80 GPU。
測試代碼鏈接:https://colab.research.google.com/drive/1b3QpfSETePo-J2TjyO6D2LgTCjVrT1lu
下表是結(jié)果摘要。在同樣情況下,將數(shù)據(jù)從 PyTorch CUDA 張量傳遞到 CUDA PyTorch 嵌入變量上是要比 SpeedTorch 更快的,但對于所有其他的傳輸類型,SpeedTorch 更快。對于轉(zhuǎn)移到 Cuda Pytorch 嵌入,或從 Cuda Pytorch 嵌入轉(zhuǎn)移的兩個步驟的總和上來說,SpeedTorch 比常規(guī) GPU 和 CPU Pinned 張量的 Pytorch 速度同樣快。
聯(lián)系客服