【機器學習排坑系列】記憶體不夠怎麼辦？

那些坊間的教學書不會跟你提到的故事。

6 min readJul 5, 2020

[注意]本文所提供的只是一個解決辦法，但並只有文中所提到的方式能解決記憶體不足的問題。
[訓練語言]：Python3

在機器學習的過程中，大部分剛入門的人會遇到的第一個大問題，或許就是在利用大量數據進行訓練時所出現的記憶體不足問題。現在的運算設備，若連同雲端運算服務一起考慮的話，記憶體大於128 GB的實在不多，但在這個記憶體大小之下，勢必無法將整筆數據一次性的導入到記憶體中，那這時候該怎處理呢？我們可以在這篇文章中做些簡單的討論。

什麼是資料流？能吃嗎？

在開始討論記憶體不足之前，我們可以先討論一下關於機器學習的訓練流程以及一些關於資料流的問題。

在一般入門的人眼中，訓練一個模型的流程大概是：

讀取資料
對資料進行預處理
建立模型，決定超參數
開始訓練並評價結果，若表現不佳則重新訓練
進行超參數微調
輸出模型

如果用一張流程圖來表達的話，就如同下圖：

在這種流程之下，資料會在「讀取資料」的步驟時就被從硬碟讀取到記憶體上，並進行後續的處理（預處理、訓練等等）。一般來說，這些資料直到訓練結束之前，都會一直在記憶體之上，直到程式結束並釋放記憶體為止。這是因為Python的設計上是直到程式結束點才會主動呼叫解構子（destructor）並釋放本身所佔用的記憶體所導致。

那麼問題來了，這些資料大概會佔用多少記憶體呢？以一個彩色（RGB）照片來說，一張解析度為224*224的照片所佔用的記憶體約為224*224*3位元組，若有一萬張這樣的照片的話，就至少會佔用約1.5十億位元組（1.5GB）的記憶體空間。

說到這裡，大家應該會發現一個問題，單單一萬張彩色照片就會佔用約1.5GB的空間，那麼在進行更大筆資料（十萬筆資料或更多）及的訓練時，記憶體的佔用程度豈不是更加的多？況且預處理以及訓練的過程還會有額外的記憶體佔用，這下記憶體怎麼可能夠用？以Google提供的Colab來說，免費使用的記憶體額度約為15GB左右，一般的電腦或是工作站最高約為128~256GB，可以想見若以相當大的資料集（以ILSVRC2012 訓練資料集壓縮檔為例，其大小為147.9GB）進行訓練，記憶體肯定是不夠用的。

讓我們來看一個簡單的例子。

import <packages>def get_data(): #read data from file
....
def pre_process(): #pre-processing 
....
def construct_model(): #construct a model
....
def train(): # define train function
....def main():
    x, y = get_data()
    img, label = pre_process()
    model = construct_model()
    train_history, trained_model = train()
#end

以上方這個程式來說，就是一個符合圖（一）所表達的訓練流程。這種流程將會在執行完所有程式碼之後，結束並程式並釋放記憶體。參考我個人的Github頁面，這種方式在ResNet-50的模型上，讀取3萬筆224*224的彩色圖片進行訓練就是極限了。下圖是執行GitHub範例的記憶體使用量（採樣到開始訓練一段時間為止）。

山窮水盡疑無路，柳暗花明又一村？

既然如此，我們該如何處理這個記憶體不夠的問題呢？在上一段我們有提到一個重點：

「這些資料直到訓練結束之前，都會一直在記憶體之上，直到程式結束並釋放記憶體為止。」

根據上面這一句話，最簡單就是將訓練過程拆成好幾個區段，一次只讀取一部分來進行訓練，並在該區段結束後將程式結束，並釋放記憶體。

說得很簡單，實際上該怎麼做呢？我們可以舉一個簡單的流程：

對於上方這種訓練流程來說，每完成一次「讀取」至「儲存訓練紀錄」的週期後，就將程式關閉，並重新啟動一次程式，並從上一次結束的進度開始訓練。這種流程的好處是：

每次都會釋放出記憶體，使得系統負載相對低。
將預處理資料寫入硬碟可以使得記憶體佔用降低。
利用tensorflow_io模組所讀取的資料將不會直接進入記憶體，而是以生成器（generator）的方式讀取，可以更近一步地降低記憶體使用。

但有個上方流程壞處是，使用者必須自行設計完整的資料流以及在程式重啟時的銜接。不過這些對於記憶體不足問題來說算是相對小事就是了。

我們來看一個使用這種訓練流程的範例。本範例是基於Python與Shell腳本來實現的一種做法。

#!/bin/bash 
train="train"
test_="test"  
for i in {0..102..1};do     
init=$i    
fin=$[$i + 1]    
python3 ResNet-50_tf_keras-v6tmp.py $init $fin $train
done 
for i in {0..33..1};do     
init=$i    
fin=$[$i + 1]    
python3 ResNet-50_tf_keras-v6tmp.py $init $fin $test_
done

程式主體可以參照這個GitHub連結。對於這種訓練方式，其記憶體使用量可以有大幅度地降低（如下圖）。

參考圖（二），我們可以看到與圖（一）的訓練流程相比，圖（三）的訓練流程具有以下優勢：

讀取以及資料預處理時的記憶體佔用較低
訓練時所使用的記憶體佔用被極大程度的降低（近120GB →約18GB）

不過，若仔細觀察程式主體，將會發現這種訓練方式的整體速度會被硬碟的讀取速度所限制。不過相對於記憶體不夠導致無法以大量資料進行訓練，這種方式所帶來的時間影響基本上是可以被接受的。基本上可以被理解成一個以空間換取時間的做法。

[補充]：還有一種是基於gc的解決方式。由於gc不是一種很推薦的方式，所以本文暫時不討論這種處理。

結語

在本文中，我們簡單的討論了關於模型訓練時記憶體不足的問題以及關於解決方法的思路，並參考了一些案例。藉由本文，我們可以兩種訓練流程在對記憶體的使用上有著非常大的差異，而這個差異是無法被忽視的。

基本上這只是眾多處理方式的一種，若讀者有更好的處理方式，也歡迎在下方進行回應。

希望這篇文章能幫助到遇到訓練時記憶體不足的問題的朋友。

【機器學習排坑系列】記憶體不夠怎麼辦？

那些坊間的教學書不會跟你提到的故事。

什麼是資料流？能吃嗎？

山窮水盡疑無路，柳暗花明又一村？

結語

Written by David Ho