N3LDG:一種輕量級自然語言處理深度學習庫

王潛升 余南 張梅山 韓子嘉 付國宏

黑龍江大學計算機科學技術學院, 哈爾濱 150080; ? 通信作者, E-mail： ghfu@hotmail.com

近年來, 深度學習方法在自然語言處理領域很多任務中的性能超越了傳統(tǒng)的統(tǒng)計機器學習方法,得到廣泛的應用[1]。在訓練階段, 深度學習模型的執(zhí)行步驟包括前向傳播、反向傳播和更新參數等,深度學習庫能方便地執(zhí)行這些步驟。Theano[2]、CNTK[3]、Caffe[4]、TensorFlow[5]和PyTorch[6]等已得到廣泛的應用[7-10]。

Theano[2]、CNTK[3]、Caffe[4]和TensorFlow[5]在訓練前靜態(tài)地定義計算圖, 在訓練時對所有實例執(zhí)行同一個計算圖。然而, 在自然語言處理任務中,構建適應所有實例的計算圖存在額外的困難, 體現(xiàn)在以下兩方面。

1)各實例的長度不一致。補零可使各實例長度一致, 然而補零操作可能影響計算結果。為了避免這個影響, 需對計算結果做裁剪。

2)實例含有結構化信息, 如句法結構。有時我們希望基于這些結構化信息動態(tài)地構建計算圖, 比如在基于句法的遞歸神經網絡中, 不同的句子實例有著不同的句法結構, 也就對應不同的計算圖。

PyTorch等深度學習庫則根據不同的實例, 動態(tài)地構建不同的計算圖。為了利用多核CPU或GPU加速計算, PyTorch要求使用者將可以批量化計算的數據手動合并為張量。比如在計算機視覺任務中,使用者須將多個圖像實例合并為一個張量后作為模型的輸入。

然而, 在自然語言處理任務中, 手動批量化合并數據存在以下額外的困難： 1)各句子實例須在補零后才能合并為張量; 2)在樹結構模型(如遞歸神經網絡)中, 須分析計算圖執(zhí)行步驟后, 將同一執(zhí)行步驟中處于不同實例上的計算過程批量化。

Looks等[11]提出一種自動批量化方法, 允許深度學習庫構建完整個計算圖后, 自動地發(fā)現(xiàn)當前可執(zhí)行的同類型計算過程, 并將其批量化執(zhí)行。為方便自然語言處理領域的研究者使用, 我們也實現(xiàn)了動態(tài)計算圖和自動批量化。如同多數深度學習庫[2-6],我們還實現(xiàn)了自動微分。N3LDG將向量視為計算的對象, 將卷積、池化等視為基于向量的各種操作,而在自然語言處理任務中, 深度學習模型的輸入通常是詞向量, 或拼接了其他特征的向量, 這樣N3LDG滿足了自然語言處理任務的要求。為提高執(zhí)行速度, N3LDG使用C++語言來實現(xiàn)。與其他深度學習庫相比, N3LDG更容易使用, 只需在項目中包含頭文件即可使用。在Apache 2.0協(xié)議下, N3LDG在https：//github.com/zhangmeishan/N3LDG發(fā)布。

1 相關工作

近年出現(xiàn)很多通用深度學習庫。Zhang等[12]提出一種自然語言處理深度學習庫LibN3L, 實現(xiàn)深度學習模型中的常見操作, 但是該庫不支持自動批量化。針對深度學習模型的計算圖的自動批量化研究尚不多見。Looks等[11]首先提出基于節(jié)點在計算圖中的深度的自動批量化方法后, Neubig等[13]認為這個方法在處理RNN模型時難以充分批量化。為緩解這個問題, Neubig等[14]提出一種將同類型節(jié)點在計算圖中的平均深度作為啟發(fā)式規(guī)則的方法, 并應用在他們的深度學習庫DyNet中。由于RNN模型在自然語言處理任務中較常用, 為了高效地訓練RNN模型, 我們仿照Neubig等[13]的方法。

多數的深度學習庫能夠利用GPU加速訓練模型[2-6,14]。Chetlur等[15]提出cuDNN庫, 高效地實現(xiàn)深度學習中的各基本操作。為了高效地分配顯存,DyNet在庫初始化時創(chuàng)建了3個顯存塊[14], 其中一個顯存塊在前向傳播中使用, 另一個在反向傳播中使用, 最后一個用于存儲參數和相關的梯度。這樣,可通過指針的加減運算, 實現(xiàn)分配和釋放顯存的操作。我們利用顯存池來高效地分配顯存, 這種做法不要求使用者預估顯存占用空間, 而是在需要時動態(tài)地向系統(tǒng)申請新的顯存塊。

2 計算圖

2.1 計算圖的引入

對于一個簡單的線性分類器y=Wx, 只需以x為自變量做線性變換, 便可得到分類結果。為說明計算圖的優(yōu)點, 我們引入更復雜的模型。圖1描述一種循環(huán)神經網絡(RNN)模型。該模型可表示為y=f(x1,x2, …,xn), 我們難以直接表示f, 因此將f分解為多個簡單的計算步驟, 每步計算結果存儲在一個中間向量中。將向量視為圖的節(jié)點, 向量之間形成有向邊, 分解后的各計算步驟和向量構成計算圖G。

圖1 一種RNN模型Fig.1 An RNN model

為實現(xiàn)計算圖, 首先定義Node類作為計算圖節(jié)點。以圖1中hi=tan h （Whhi-1+zi）為例, 為了計算hi,Node類需包含以下信息： 1)前向傳播的計算方法;2)本節(jié)點向量(hi); 3)各輸入向量(hi-1,zi); 4)參數(Wh)。當給定各輸入向量x1,x2, …,xn時, 某節(jié)點即可執(zhí)行前向傳播過程, 求得本節(jié)點向量y=f(x1,x2, …,xn), 稱該節(jié)點為可執(zhí)行節(jié)點。y又可作為其子節(jié)點的輸入向量, 使子節(jié)點成為可執(zhí)行節(jié)點。若某節(jié)點不含輸入向量(如圖1中xi所在節(jié)點), 則該節(jié)點成為計算圖G的初始可執(zhí)行節(jié)點。這樣, 以初始可執(zhí)行節(jié)點為始, 重復執(zhí)行前向傳播過程, 直至計算圖中所有節(jié)點都被執(zhí)行, 即完成模型的前向傳播過程。圖 1 描述了這個過程。

為執(zhí)行反向傳播, Node類還須包含以下信息： 1)反向傳播的計算方法; 2)損失函數L對y的導數;3)L對各輸入向量的導數4)L對各參數的導數我們在Node類中定義前向傳播和反向傳播的接口, 在其各個子類中實現(xiàn)這兩個接口。我們實現(xiàn)了常用的節(jié)點類型,包括tanh, concat和線性變換等, 列舉在表 1 中。使用者也可自己定義新的節(jié)點類型, 實現(xiàn)前向傳播和反向傳播。

表1 N3LDG中的常用節(jié)點類型Table 1 Commonly used node types in N3LDG

2.2 自動批量化

計算圖中往往有多個可執(zhí)行節(jié)點。為了提高執(zhí)行速度, 需要批量化地執(zhí)行同類型的計算過程。具體而言, 有兩類計算過程可批量化執(zhí)行： 1)共享參數的同類型計算, 如y1=Wx1+b和y2=Wx2+b; 2)不含參數矩陣的同類型計算, 如y1=tanh(x1)和y2=tanh(x2)。

N3LDG自動發(fā)現(xiàn)當前可執(zhí)行節(jié)點, 并批量化執(zhí)行同類型的計算過程。當這些節(jié)點被執(zhí)行完后, 即從計算圖中移除, 此時可得新的可執(zhí)行節(jié)點集合。這樣, 我們總能得到當前可執(zhí)行節(jié)點的集合, 直到計算圖執(zhí)行完畢。以圖1中RNN模型為例, 執(zhí)行步驟如下。

1)[x1x2…]=[emb(大家)emb(好)… emb(！)]

2)[z1z2…]=Wx[x1x2…]+[b b…b]

3)[h1]=tanh([z1)

4)[h2]=tanh(Wh[h1]+[z2])

5)[h3]=tanh(Wh[h2h2′]+[z3])

8)[p p′ ]=[pool(h1,h2,h3)pool

9)[y y′ ]=Wp[p p′]

3 CPU計算

Eigen是通用的C++線性代數計算庫[16], 因此我們使用Eigen實現(xiàn)CPU上的線性代數計算。由于CPU能高效地處理內存中連續(xù)存放的向量, 所以N3LDG對常用的計算過程做了優(yōu)化。具體的步驟如下： 1)計算前, 將各節(jié)點中的輸入向量合并為一個矩陣, 將矩陣與多個向量的乘法運算轉換為矩陣與矩陣的乘法運算; 2)執(zhí)行矩陣和矩陣的乘法運算,得到結果矩陣; 3)將該結果矩陣拆分后, 賦值給各節(jié)點的向量。

我們以y1=tanh(Wx1+b),y2=tanh(Wx2+b)...yn=tanh(Wxn+b)為例, 首先將x1,x2,…,xn合并為矩陣［x1x2…xn］, 記為X, 將同一個向量b擴展為n列矩陣［b b...b］, 記為B。將各向量拷貝至連續(xù)的內存區(qū)域中, 然后執(zhí)行Y=tanh(WX+B), 計算完成后,將矩陣Y拆分拷貝至各節(jié)點的向量。

4 GPU計算

cuBLAS是英偉達發(fā)布的CUDA線性代數計算庫, 我們使用cuBLAS實現(xiàn)GPU上的線性代數計算,并編寫kernel函數實現(xiàn)其余計算過程。為充分利用GPU的并行計算能力, 我們并行執(zhí)行所有批量化之后的計算過程。我們的實現(xiàn)不依賴cuDNN[15], 使用者無需安裝cuDNN。

我們發(fā)現(xiàn)GPU中有兩類操作存在性能瓶頸： 1)顯存分配與釋放; 2)顯存和內存間的I/O。當動態(tài)構建計算圖時, 參與前向傳播和反向傳播計算過程的各向量地址也隨之動態(tài)地變化, 計算前, 須將這些信息傳輸到顯存, 這會頻繁涉及上述兩類操作。我們通過以下方法來緩解性能瓶頸。

4.1 顯存分配與釋放

在實驗中測量顯存的分配與釋放時間, 發(fā)現(xiàn)它們占總訓練時間相當大的比例, 成為性能瓶頸。通過專用模塊(顯存池), 持有并管理空閑的顯存塊,當不持有合適的空閑塊時, 才向系統(tǒng)申請顯存塊,從而減少向系統(tǒng)分配與釋放顯存的次數。英偉達實現(xiàn)了顯存池庫cnmem, 在https：//github.com/NVIDIA/cnmem發(fā)布。Knowlton等[17]提出伙伴系統(tǒng), 用于快速分配存儲空間。受伙伴系統(tǒng)啟發(fā), 我們也實現(xiàn)了顯存池。

4.2 I/O

對于同樣大小的數據, 只調用一次庫函數, 將其傳輸至顯存, 顯著地快于分成多次傳輸[18]。因此,對需傳輸至顯存的多個數據, 我們在內存中將其連續(xù)存放后, 再調用一次庫函數傳至顯存。比如, 批量化執(zhí)行y1=tanh(x1),y2=tanh(x2)...yn=tanh(xn)時, 需將x1,x2, …,xn的地址傳輸至顯存, 我們在內存中連續(xù)存放x1,x2, …,xn的地址后, 調用一次庫函數, 將這些地址傳輸至顯存。與調用多次庫函數而分別傳輸它們的地址相比, 我們的方法顯著地減少了顯存與內存間的I/O次數。

5 實驗

我們通過一個 5 分類情感分類任務, 在 3 個模型上做基準測試： 1)卷積神經網絡(CNN); 2)雙向長短時記憶網絡(Bi-LSTM); 3)樹結構長短時記憶網絡(Tree-LSTM)[19]。以上模型的詞向量和隱層的維度都設置為200。訓練數據包括8544個句子實例,共163563個詞(包括標點符號), 測試代碼在https：//github.com/chncwang/n3ldg-benchmark發(fā)布。記錄訓練一輪epoch的時長, 包括： 1)構建計算圖; 2)規(guī)劃執(zhí)行步驟; 3)前向傳播; 4)反向傳播; 5)更新參數。實驗中, CPU的型號是Intel (R)Core (TM)i7-6800K CPU @ 3.40 GHz, GPU型號是GeForce GTX 1080 Ti。我們還用PyTorch實現(xiàn)一致的模型結構,并在 3 個模型上實現(xiàn)手動批量化, 以便與N3LDG對比訓練速度。在N3LDG中, 我們未對LSTM做特別的優(yōu)化, 為公平對比, 用PyTorch以同樣方式實現(xiàn)LSTM。

首先在單線程CPU上, 對N3LDG和PyTorch做基準測試, 測試結果見表2。

表2 顯示, 在所有設置下, N3LDG在單線程CPU上的訓練速度高于PyTorch。當我們訓練CNN時,N3LDG訓練速度達到PyTorch的9.40～42.47倍, 訓練Bi-LSTM時達到PyTorch的4.43～9.71倍, 訓練Tree-LSTM時達到PyTorch的1.28～3.10倍。這表明我們構建計算圖、自動批量化和CPU計算過程是高效的。

表2 單線程CPU上N3LDG和PyTorch的基準測試Table 2 Benchmarks of N3LDG and PyTorch on single thread CPU

我們在GPU上對N3LDG、不使用cuDNN的PyTorch (稱為PyTorch CUDA)以及使用cuDNN的PyTorch (PyTorch cuDNN)做了基準測試, 測試結果見表 3。

表3 顯示, 在訓練CNN和Tree-LSTM時, N3LDG在GPU上的訓練速度高于PyTorch CUDA和PyTorch cuDNN。在訓練CNN時, N3LDG的訓練速度達到PyTorch CUDA的3.40～18.38倍, PyTorch cuDNN的3.10～8.74倍, 訓練Tree-LSTM時速度達到PyTorch CUDA的1.78～3.03倍, PyTorch cuDNN的1.67～2.79倍。訓練Bi-LSTM模型時, N3LDG在較大的minibatch下有優(yōu)勢。當mini-batch=1時, N3LDG的訓練速度低于PyTorch, 是PyTorch CUDA的77.46%,PyTorch cuDNN的80.18%。當mini-batch=16時,N3LDG的訓練速度與PyTorch幾乎相同。當minibatch=256時, N3LDG的訓練速度達到PyTorch CUDA的1.71倍, PyTorch cuDNN的1.77倍。總體而言, 我們構建計算圖、自動批量化和GPU計算過程是高效的。

為了分析自動批量化對訓練速度的影響, 以Bi-LSTM (MB=256)為例, 分別在單線程CPU和GPU上測試是否做自動批量化時各步驟的時長。實驗結果見圖 2。

圖2顯示, 自動批量化顯著地提升了訓練速度。在單線程CPU上提升4.76倍, 在GPU上提升52.27倍。提速主要來自前向傳播和反向傳播。

我們猜測在單線程CPU上提速的部分原因在于合并了矩陣與向量的乘法, 比如將y1=Wx1和y2=Wx2轉換為[y1y2]=W[x1x2]后, 計算速度更快。我們還對比了在相同設置下, 一輪epoch中矩陣乘法的總執(zhí)行時間、執(zhí)行次數及平均執(zhí)行時間, 結果見表 4。

表3 GPU上N3LDG, PyTorch CUDA和PyTorch cuDNN的基準測試Table 3 Benchmarks of N3LDG, PyTorch CUDA and PyTorch cuDNN on GPU

圖2 做自動批量化與否時各訓練階段時長Fig.2 Time in per training stage when auto-batch enabled or not

表4 顯示自動批量化顯著地提升了單線程CPU矩陣乘法的執(zhí)行速度, 提升幅度為9.28倍。與不做批量化相比, 盡管自動批量化時平均執(zhí)行時間更長,但執(zhí)行次數僅為0.98%。

為分析自動批量化對CUDA核函數執(zhí)行速度的影響, 我們對比了在相同設置下, 一輪epoch中核函數的總執(zhí)行時間、執(zhí)行次數及平均執(zhí)行時間, 結果見表 5。

表5 顯示, 自動批量化顯著提升了CUDA核函數的執(zhí)行速度, 提升幅度為70.52倍。與不做批量化相比, 執(zhí)行次數僅為1.28%。值得注意的是, 平均執(zhí)行時間只是不做批量化時的1.11倍, 表明自動批量化充分利用了GPU的并行計算能力。

表4 單線程CPU矩陣乘法的總執(zhí)行時間、執(zhí)行次數和平均執(zhí)行時間Table 4 Total execution duration, times and average duration of matrix multiplication

表5 CUDA核函數的總執(zhí)行時間、執(zhí)行次數和平均執(zhí)行時間Table 5 Total execution duration, times and average duration of CUDA kernel functions

圖3 無顯存池、cnmem和我們的顯存池訓練時間對比Fig.3 Comparison of training time among the absence of the memory pool, cnmem and ours

為測試顯存池的有效性, 我們在訓練Bi-LSTM時, 分別統(tǒng)計一輪epoch中, 不使用顯存池、使用cnmem以及使用我們的顯存池時的訓練時間以及分配與釋放顯存的時長, 實驗結果見圖 3。

圖3 顯示, 當不使用顯存池時, 顯存的分配與釋放占訓練時間的56.87%～74.72%, 成為性能瓶頸,而顯存池顯著降低時長。使用我們的顯存池時, 分配與釋放顯存的速度是使用cnmem時的5.11～37.11倍, 訓練速度是無顯存池時的3.19～3.37倍, 使用cnmem時的1.13～2.63倍, 表明我們的顯存池是高效的。

6 結論

為方便在自然語言處理任務中應用深度學習,移除手動批量化過程, 本文提出一種輕量級自然語言處理深度學習庫N3LDG。我們仿照Neubig等[13]的方法, 實現(xiàn)自動批量化, 并在CPU和GPU上都高效實現(xiàn)常見的深度學習計算過程。實驗表明, 自動批量化顯著提高了CPU和GPU上的執(zhí)行速度。我們的庫在CNN, Bi-LSTM和Tree-LSTM模型中的CPU性能以及在CNN和Tree-LSTM模型中的GPU性能都優(yōu)于PyTorch。作為一種輕量級的庫, 我們開發(fā)的N3LDG為自然語言處理領域的研究者提供了新的選擇。