Categories
程式開發

如何優化TensorFlow Lite 運行時內存?


由於資源有限,在移動和嵌入式設備上運行推理很有挑戰性。人們必須在嚴格的能耗要求下使用有限的硬件。在本文中,我們希望展示TensorFlow Lite(TFLite)內存使用方面的改進,這些改進使其更適合在邊緣設備運行推理。

中間張量

通常情況下,神經網絡可以被認為是計算圖,由運算符(如CONV_2DFULLY_CONNECTED)和包含中間計算結果的張量(稱為中間張量)組成。這些中間張量通常是預先分配的,以減少推理等待時間,而代價是存儲空間。然而,如果天真地實現了這個代價,那麼在資源受限的環境中就不能掉以輕心:它可能會佔用大量的空間,有時甚至比模型本身還要大好幾倍。例如,MobileNet V2 中的中間張量佔用了26MB 內存(見圖1),大約是模型本身的兩倍。

如何優化TensorFlow Lite 運行時內存? 1

圖1:MobileNet V2 的中間張量(上)及其在2D 存儲空間(下)上的映射。如果每個中間張量使用專用的內存緩衝區(用56 種不同的顏色表示),它們將佔用約26MB 的運行時內存

好消息是,由於數據相關性分析,這些中間張量不必在內存中共存。這允許我們可以重用中間張量的內存緩衝區,並減少推理機的總內存佔用。如果網絡具有簡單鏈的形狀,那麼兩個大的內存緩衝區就足夠了,因為它們可以在整個網絡中來回互換。然而,對於構成複雜圖的任意網絡,這種NP 完全(NP-Complete,縮寫為NP-C 或NPC)資源分配問題需要一個良好的近似算法。

我們為這個問題設計了許多不同的近似算法,它們的執行方式都取決於神經網絡和內存緩衝區的屬性,但它們都使用一個共同點:張量使用記錄。一個中間張量的張量使用記錄是一個輔助數據結構,它包含關於張量有多大,以及在網絡的給定執行計劃中第一次和最後一次使用的時間信息。在這些記錄的幫助下,內存管理器能夠在網絡執行的任何時刻計算中間張量的使用,並優化其運行時內存以獲得盡可能小的佔用空間。

共享內存緩衝區對象

TFLite GPU OpenGL 後端,我們為這些中間張量採用GL 紋理。這些都有一些有趣的限制:(a)紋理的大小在創建之後便無法修改,並且(b)在給定的時間只有一個著色器程序獲得對紋理對象的獨占訪問權。在這種共享內存緩衝區對像模式中,目標是最小化對像池中所有創建的共享內存緩衝區對象的大小之和。這種優化類似於眾所周知的寄存器分配問題,只不過由於每個對象的大小不同,這種優化要復雜得多。

根據前面提到的張量使用記錄,我們設計了5 種不同的算法,如表1 所示。除了最小成本流之外,它們都是貪心算法(Greedy algorithm),每種算法使用了不同的啟發式算法,但仍然非常接近或接近理論下限。根據網絡拓撲的不同,一些算法比其他算法執行得更好,但通常情況下,GREEDY_BY_SIZE_IMPROVEDGREEDY_BY_BREADTH產生的對象分配佔用的內存最少。

如何優化TensorFlow Lite 運行時內存? 2

表1:共享對象策略的內存佔用(以MB 為單位;最佳結果以綠色高亮顯示)。前5 行是我們的策略,後2 行作為基線(下限表示可能無法實現的最佳數值的近似值,Naive 表示每個中間張量分配自己的內存緩衝區時可能出現的最差數值)

回到前面的示例,GREEDY_BY_BREADTH在MobileNet V2 上執行得最好,它利用了每個運算符的寬度,即運算符配置文件中所有張量的總和。圖2,特別是與圖1 相比時,突出顯示了使用智能內存管理器時,可以獲得多大的收益。

如何優化TensorFlow Lite 運行時內存? 3

圖2:MobileNet V2 的中間張量(上)及其在2D 存儲空間(下)上的映射。如果中間張量共享內存緩衝區(用4 種不同的顏色表示),它們只佔用約7MB 的運行時內存

內存偏移量計算

對於在CPU 上運行的TFLite,適用於GL 紋理的內存緩衝區屬性並不適用。因此,更常見的做法是預先分配一個巨大的內存空間,並讓所有訪問它的讀取器和寫入器共享它,這些讀取器和寫入器通過給定的偏移量訪問它,而不會干擾其他讀取和寫入。這種內存偏移量計算方法的目標是最大程度地減少內存空間

我們已經為這個優化問題設計了3 種不同的算法,並探索了先前的研究工作(Sekiyama 等人在2018 年提出的Strip Packing)。與共享對象方法類似,有些算法的性能優於其他算法,具體取決於網絡,如表2 所示。從這個調查中我們可以看出,偏移量計算方法比一般的共享對象方法佔用的空間更小,因此,如果可行的話,我們應該選擇前者而非後者。

如何優化TensorFlow Lite 運行時內存? 4

表2:偏移量計算策略的內存佔用(以MB 為單位;最佳結果以綠色高亮顯示)。前3 行是我們的策略,接下來的1 行是先前的研究工作,最後2 行作為基線(下限表示可能無法實現的最佳數字的近似值,而Naive 表示每個中間張量分配自己的內存緩衝區可能出現的最差數值)

這些針對CPU 和GPU 的內存優化,默認情況下已經與最近幾個穩定的TFLite 版本一起提供,並且在支持諸如MobileBERT等更高要求、最先進的模型方面證明了它們的價值。你可以通過直接查看GPU 實現CPU 實現來找到有關實現的更多詳細信息。

原文鏈接:

https://blog.tensorflow.org/2020/10/optimizing-tensorflow-lite-runtime.html