「深度解析」AI訓練之數據緩存

深度學習或者AI的出現，改變了咱們以往的解決問題的編程方式，再也不是代碼上直觀的表達。

舉一個簡單的例子，咱們如何識別一個數字（圖片）是數字9呢？很是直觀的方法就是上面有一小圓圈，下面有一個豎線。可是若是寫的傾斜了一些呢？若是上面的圈沒有閉合呢？若是豎線彎曲了呢？感受咱們平常的程序判斷（switch）沒法收斂，咱們只能用一種可以自我演進的方式來認識這個「看起來像9」的數字，而這也正是咱們大腦的學習行爲，咱們第一個看到這個數字的時候，被告知這是9，那麼圖片就有了一個標籤；下次再看到相似的，仍是屬於標籤9，見多識廣，最後見到一個也許寫得更加不像的咱們也可以識別出是9，這個過程正是由咱們大腦中的上千億的神經細胞長時間的學習結果。

人類大腦的真正運行方式，依舊是神祕所在，但從這個過程當中咱們發展出了神經網絡算法，能夠從已有的知識中進行學習。勤能補拙，既然算法不如人腦，就經過學習大量的資料來加快學習的進程。MNIST 數據有 60,000張手寫數字圖片，ImageNet 數據有接近1500萬張圖片，Youtube-8M的視頻數據有數TB，Google的 Open Image dataset, 僅僅是在 Open Images Challenge中使用的數據集就達到了18TB。

AI中有三大核心：算法，算力，數據（存儲）。 算法自有成熟的框架，由數學科學家去解決；計算能力由CPU甚至GPU去解決。面對如此大量的數據，一臺機器的內存、硬盤去承載基本不太可能，而對於CPU/GPU計算能力強悍的組件，頻繁的去遠端獲取數據等待IO又是資源的浪費。有沒有既能知足數據距離計算近、又能承載大量數據的方案呢？緩存是銀彈！後面的主要篇幅從論文解析的角度來逐步闡述，論文來自Fast20 Quiver: An Informed Storage Cache for Deep Learning。

後續的討論中，有個比較重要的概念，就是mini-batch, 若是沒有實戰的經歷過，不是很容易理解這個概念。

深度學習的優化算法，本質就是梯度降低。每次的參數更新有兩種方式。

第一種，遍歷所有數據集計算一次損失函數，而後計算函數對各個參數的梯度，更新梯度。這種方法每更新一次參數都要把數據集裏的全部的樣本數據遍歷一遍，計算量開銷大，計算速度慢，不支持在線學習，這稱爲Batch gradient descent（BGD），批梯度降低。

另外一種，每訓練一個數據就算一下損失函數，而後求梯度更新參數，這個稱爲隨機梯度降低，stochastic gradient descent。這個方法速度比較快，可是收斂性能不太好，可能在最優勢附近波動，達不到最優勢。兩次參數的更新也有可能互相抵消掉，形成目標函數震盪的比較劇烈。

爲了克服兩種方法的缺點，如今通常採用的是一種折中手段，mini-batch gradient decent，小批的梯度降低，這種方法把數據分爲若干個批，按批來更新參數，這樣，一個批中的一組數據共同決定了本次梯度的方向，降低起來就不容易跑偏，減小了隨機性。

用一個示意圖表示以下：

藍色：爲 batch 梯度降低，即 mini batch size ＝ m，紫色：爲 stochastic 梯度降低，即 mini batch size ＝ 1，綠色：爲 mini batch 梯度降低，即 1 < mini batch size < m。

如下圖爲例，執行了3輪訓練（epoch）,每輪裏面定義mini-bach size=5, 其中數據集爲1-20個數字，咱們看到經過torch.DataLoader, 每次得到了5個數據(batch x)。

01 深度學習訓練的基本知識

深度學習訓練任務(Deep Learning Training DLT)會將訓練數據做爲輸入，從千絲萬縷的線索中經過學習並獲得一個輸出模型來表明訓練數據。

爲了實現訓練，DLT會使用一個較小的隨機樣本（mini-batch,一般是32到512個），並利用SGD來慢慢的學習各類參數進而提升準確率。

訓練數據：一般咱們能夠認爲是一個列表，列表中的每個元素都是一個二元組<input,label>, input多是一張圖片或者一段語音，而label則表明着input的語義，而這也正是深度學習網絡所須要學習的並可以正確區分input的目標。例如ImageNet的所有數據集大概有150萬張圖片，每張圖片在200KB左右。

爲了可以以隨機方式訪問訓練數據，DLT框架會使用索引序列來遍歷數據。假設訓練數據有100萬個文件，那麼會維護一個包含每個文件索引的列表，並對它進行隨機的排列，隨後根據mini-batch的數據量向後端存儲得到數據，當所有的數據都完整遍歷訓練一次，一個epoch完成。對於下一個epoch, 再次對索引進行隨機排列，重複上面的過程。一個典型的DLT任務會運行不少輪訓練，例如50-200。

數據轉換：從存儲得到原始數據會被DLT框架進行轉換，例如彩色圖片變成黑白圖片，同時將圖片轉換爲像素數矩陣等等。固然這部分工做一般由CPU來完成。

多任務：由於DLT任務是一個試錯的過程，因此實際運行過程當中，用戶老是會使用不一樣的參數來同時運行不一樣的任務，全部的這些任務都會訪問相同的完整數據集，不一樣的就是以不一樣的隨機順序來進行訪問。

02 深度學習的IO特色

咱們從DLT任務I/O訪問的角度看來列舉一下它的主要特色：

可共享性：在DLT訓練任務中，不管是一個訓練任務自身，仍是多個訓練任務之間，都存在很大程度的I/O重疊性。在一個任務內，它會針對同一個數據集進行屢次的遍歷（例如多個epoch），因此若是可以在第一個epoch的時候就對數據進行緩存，會大幅提高後續訓練的效率。更重要的是，這種共享性甚至能夠擴展到多任務之間，例如針對同一份訓練數據集，配置不一樣的參數，利用同一個訓練模型運行多個不一樣的任務。這些任務可能運行在不一樣的機器上，可是訪問的都是相同的底層數據集。

隨機訪問：因爲數據的可共享性，這使得DLT具備很是的緩存友好性，但只有在所有數據可以被完整緩存的狀況下才有效果，不然，DLT隨機訪問數據的方式又使得部分數據緩存很容易被穿透。例如只可以緩存20%的數據，那麼這些數據立刻就會被後續的隨機訪問刷掉。

部分數據緩存對於DLT來講很重要，由於訓練數據一般已經足夠大，而且會愈來愈多，例如前文提到過即便只是ImageNet這樣的百萬級規模的數據集，整體也已經達到了數TB的大小。

可替換性：從I/O的角度來講，一個訓練任務(epoch)主要關注如下兩點便可：a)每個訓練數據必須且僅被訪問一次；b)而對於每次的mini-batch,必須是隨機的序列。有趣的是，一組數據的精確順序並不會對訓練任務的準確或者精確性產生影響，這就意味着I/O是能夠被替換的。對於特定若干文件的訪問，DTL任務能夠替換爲一組其餘的隨機的且沒有被訪問過的數據。從緩存的角度來講，這是一個獨特的特性來提高緩存的命中率，即便緩存只能承載20%的數據，也能夠在訪問一個不存在於緩存中的數據，經過替換的方式返回一個存在的內容，同時並無破壞隨機以及惟一性的訓練要求。

可預測性：由於每個mini-batch的運行時間，能夠事先得到，這樣就能夠用於評估一個訓練任務對I/O性能的敏感性，進而能夠進行策略調整以可以使那麼I/O敏感的任務從緩存獲益。

03 緩存的設計

總結起來深度學習的特色：

1. 須要的數據量大
2. 多臺機器多個訓練並行
3. 每一個訓練要運行屢次
4. 在每次訓練中，全部的數據須要遍歷一遍
5. 針對不一樣的訓練參數，以及在不一樣的機器上運行的訓練任務，數據集相對保持固定

針對以上的特色，當咱們考慮緩存的時候，不由會有以下的疑問：緩存畢竟容量有限，穿透如何處理？緩存的過時置換策略是如何的？當不一樣的用戶訪問不一樣的數據，安全性如何保證？等等。

Quiver、分佈式緩存，經過與DLT框架深度整合，緩存客戶端集成到訓練任務的IO過程當中，進而爲緩存服務端提供更多的策略信息。

系統結構

以公有云虛擬機環境舉例，每個GPU VM帶有一個本地SSD硬盤，每個DLT job會運行在本身的容器內，這樣即便是多用戶運行，也是在一個隔離的環境內。

每一個用戶的數據存儲在各自帳號的雲存儲內，這樣保證了隱私以及訪問權限。經過分佈式緩存，即便訓練任務因爲調度等緣由在各個宿主之間切換，緩存數據依舊是可以提升訓練效率。

數據安全

Quiver的設計是一個共享式的分佈緩存，不管是不一樣的任務，仍是不一樣的用戶之間，在共享的模式下如何保證數據的安全就是一個重要因素。Quiver保證了用戶只能看到他有權限訪問的數據，但這樣彷佛又與緩存的重用產生了衝突。若是針對某一個數據集，例如ImageNet, 兩個不一樣用戶分別在各自的存儲帳號內各自保存了一份，那麼邏輯上來說，緩存要分別爲每一個用戶各自緩存一份。這將致使緩存的效率下降，Quiver經過內容尋址(content-addressed)的方式來解決重用與隔離的問題。

內容尋址緩存

對於緩存，基本的行爲就是經過一個<key, value>的映射關係，在咱們經過key查詢時，可以快速的返回所對應的value。在Quiver中，緩存並非利用文件名以及偏移量在做爲緩存的關鍵字，而是利用緩存內容的hashes。緩存內容的粒度是由具體的DLT任務決定的，多是一張圖片，不管是對它的插入仍是尋址，都以它的hash(例如SHA1)來惟必定位。用hash來定位的好處是對於一個相同內容的文件，無論它來自於何處以及文件名是否相同，在緩存中都僅須要保留一份便可，這樣也就可以達到即便在不一樣的用戶之間也可以共享目的。

同時爲了保證數據的隔離性，Quiver利用摘要索引來訪問訓練數據，對於每一份數據，摘要索引將包含<content_hash: file_location>， 所以，在多用戶擁有相同內容的數據集時，由於數據是存在在各自的存儲系統內，每一個用戶將擁有不一樣的file_location，可是全部的content_hash是相同的。

緩存服務器

利用本地SSD做爲介質的KV存儲，經過一致性Hash的方式將key space分佈在多個緩存服務器上。

緩存管理(Cache Manager)

因爲Quiver是分佈式緩存，那麼針對全部的緩存服務器，緩存的插入、清理須要一個協調者Cache manager。

Cache manager同時會評估每個計算任務從緩存的受益狀況，主要經過讓緩存服務器針對訓練任務所須要的若干mini-batch數據作cache misses, 而後與其餘的緩存命中的訓練人耗時機型對比，進而對緩存進行優先級調整。

緩存客戶端

緩存客戶端做爲訓練任務的一部分，經過干預DLT框架，例如PyTorch等的接口層來訪問訓練數據。在PyTorch中，DataSet會用來遍歷全部的訓練數據，而且內部維護一個隨機的文件索引列表，其中Next的接口就能夠用來得到下一個mini-batch數據。Quiver經過調整這個接口，利用一個摘要文件，當上層訪問一組文件時，它會先對緩存進行數據的訪問。

客戶端會將訓練任務的一些信息反饋給Cache Manager,例如每個mini-batch的訓練時間，Cache Manager能夠據此來優化緩存的策略。

替換命中率

在常規的緩存中，若是一個mini-batch包含了512個文件，那麼Dataset會提供512個文件索引用來從後端存儲得到文件內容，若是這其中只有部分緩存命中，那麼將依然存在遠程的I/O操做。在Quiver中，會從Cache中加載更多的（例如10倍的mini-batch數量）數據，而只要其中有512個數據可以被命中，那麼就返回給上層訓練任務，這樣訓練任務就不會被Cache miss阻塞。同時Quiver會標記Cache miss的數據爲pending狀態，周而復始，直到數據被遍歷了一遍，這時將重頭來過僅僅去關注以前pending的數據。

假設目前只有10%的數據在緩存中，爲了簡單起見，咱們能夠認爲就是連續的原始數據的10%， 由於DLT任務會隨機的查找數據，因此每個長度爲k的mini-batch序列，緩存的命中率應該爲k/10, 所以若是咱們查找一個長度爲10*k的序列，那麼正好可以命中得到mini-batch所須要的數據。當下一輪查找pending數據的時候，另外的10%的數據可能已經在緩存中了，這也意味着可以1/9的命中率。須要注意的是，在多個任務的訓練中，這依舊適用，所以多個訓練任務儘管每一個都訪問隨機的訓練數據，從總體來看，他們能夠作到以全緩存命中的方式來運行。

訓練準確性

因爲上述的I/O可替換性，咱們有理由懷疑最終訓練結果的準確性。這裏借用原文的數據來講明。

04 緩存的管理

在以前的描述中，當只有部分數據被緩存時，Quiver會在一個epoch的訓練過程當中，再次遍歷文件索引。爲了能在這後續的遍歷中得到更好的命中率，另外一部分數據必須被pre-fetch到緩存中。

Quiver經過緩存整個數據集的2個chunks來解決這個問題。首先數據集的摘要索引文件會被分紅固定數據的chunks,例如每一個chunk包含10%的數據，同時每一個chunk表明着striped partition。例如咱們定義數據集中連續的10%爲一個partition, 每一個partition被分紅10個stripe units. 這樣每一個chunk將包含全部的partition中的一個unit。這樣當訓練任務操做第一個chunk的過程當中，第二個chunk將被加載到緩存中，因此當部分訓練任務完成第一次遍歷開始第二次的時候，數據已經在緩存內，訓練任務以遞進的方式運行。

這裏面潛在的一個問題就是何時將第一個chunk置換出去。若是置換太快，部分任務尚未完成將致使緩存失效，若是保留太長時間，那麼第三個chunk將沒法加載進來。在Quiver中，當第二個chunk被加載到緩存後，第一個chunk會被標記爲能夠清除，同時新的任務能夠從第二個chunk中得到命中的數據。原有的任務依舊利用第一個chunk來運行，當全部的任務都已經遍歷了第一個chunk數據，這部分數據纔會真的從緩存中清除，同時第三部分數據將開始加載。

在上述的過程當中，若是某一個訓練任務相比於其餘的要慢不少，那麼將致使前一個chunk遲遲不能釋放，一般來講，在同一個訓練模型的多個任務中，每一個任務的訓練時間基本是相同的，但沒法避免在多個不一樣的訓練模型訓練同一個數據集的場景。不過若是一個任務明顯的耗時很長，那麼將意味着每個mini-batch在GPU上的訓練時間都很長，也就是它對I/O的性能沒那麼敏感，因此緩存的不命中並不會影響多少這個訓練的效率，所以 Quiver會設定一個臨界值來強制第一個chunk失效並清除。

05 緩存效果

論文做者經過以下的配置環境來進行效果的對比，從實際數據來看，訓練性能確實有較大的提升。

Timeline of Mini-batches

吞吐提高

06 結論

深度學習場景中，更多的注意力放在了提升計算以及網絡的性能上，而對於存儲，則是利用現有的方案來解決，例如提早手動將數據加載到離GPU較近的SSD上。論文做者經過Quiver，提供了自動化的手段來消除存儲的瓶頸。固然沒法避免對訓練框架的侵入性，但也正是由於如此，Quiver才能感知到訓練I/O的特色，進而達到即便緩存只能承載部分數據也能夠大幅調高緩存利用率的效果。