由於資源有限，在移動和嵌入式設備上運行推理很有挑戰性。人們必須在嚴格的能耗要求下使用有限的硬件。在本文中，我們希望展示 TensorFlow Lite（TFLite）內存使用方面的改進，這些改進使其更適合在邊緣設備運行推理。
中間張量
通常情況下，神經網絡可以被認爲是計算圖，由運算符（如CONV_2D或FULLY_CONNECTED）和包含中間計算結果的張量（稱爲中間張量）組成。這些中間張量通常是預先分配的，以減少推理等待時間，而代價是存儲空間。然而，如果天真地實現了這個代價，那麼在資源受限的環境中就不能掉以輕心：它可能會佔用大量的空間，有時甚至比模型本身還要大好幾倍。例如，MobileNet V2 中的中間張量佔用了 26MB 內存（見圖 1），大約是模型本身的兩倍。
好消息是，由於數據相關性分析，這些中間張量不必在內存中共存。這允許我們可以重用中間張量的內存緩衝區，並減少推理機的總內存佔用。如果網絡具有簡單鏈的形狀，那麼兩個大的內存緩衝區就足夠了，因爲它們可以在整個網絡中來回互換。然而，對於構成複雜圖的任意網絡，這種NP 完全（NP-Complete，縮寫爲 NP-C 或 NPC）資源分配問題需要一個良好的近似算法。
我們爲這個問題設計了許多不同的近似算法，它們的執行方式都取決於神經網絡和內存緩衝區的屬性，但它們都使用一個共同點：張量使用記錄。一箇中間張量的張量使用記錄是一個輔助數據結構，它包含關於張量有多大，以及在網絡的給定執行計劃中第一次和最後一次使用的時間信息。在這些記錄的幫助下，內存管理器能夠在網絡執行的任何時刻計算中間張量的使用，並優化其運行時內存以獲得儘可能小的佔用空間。
共享內存緩衝區對象
在TFLite GPU OpenGL 後端，我們爲這些中間張量採用 GL 紋理。這些都有一些有趣的限制：（a）紋理的大小在創建之後便無法修改，並且（b）在給定的時間只有一個着色器程序獲得對紋理對象的獨佔訪問權。在這種共享內存緩衝區對象模式中，目標是最小化對象池中所有創建的共享內存緩衝區對象的大小之和。這種優化類似於衆所周知的寄存器分配問題，只不過由於每個對象的大小不同，這種優化要複雜得多。
根據前面提到的張量使用記錄，我們設計了 5 種不同的算法，如表 1 所示。除了最小成本流之外，它們都是貪心算法（Greedy algorithm），每種算法使用了不同的啓發式算法，但仍然非常接近或接近理論下限。根據網絡拓撲的不同，一些算法比其他算法執行得更好，但通常情況下，GREEDY_BY_SIZE_IMPROVED和GREEDY_BY_BREADTH產生的對象分配佔用的內存最少。
回到前面的示例，GREEDY_BY_BREADTH在 MobileNet V2 上執行得最好，它利用了每個運算符的寬度，即運算符配置文件中所有張量的總和。圖 2，特別是與圖 1 相比時，突出顯示了使用智能內存管理器時，可以獲得多大的收益。
內存偏移量計算
對於在 CPU 上運行的 TFLite，適用於 GL 紋理的內存緩衝區屬性並不適用。因此，更常見的做法是預先分配一個巨大的內存空間，並讓所有訪問它的讀取器和寫入器共享它，這些讀取器和寫入器通過給定的偏移量訪問它，而不會干擾其他讀取和寫入。這種內存偏移量計算方法的目標是最大程度地減少內存空間。
我們已經爲這個優化問題設計了 3 種不同的算法，並探索了先前的研究工作（Sekiyama 等人在 2018 年提出的 Strip Packing）。與共享對象方法類似，有些算法的性能優於其他算法，具體取決於網絡，如表 2 所示。從這個調查中我們可以看出，偏移量計算方法比一般的共享對象方法佔用的空間更小，因此，如果可行的話，我們應該選擇前者而非後者。
這些針對 CPU 和 GPU 的內存優化，默認情況下已經與最近幾個穩定的 TFLite 版本一起提供，並且在支持諸如MobileBERT等更高要求、最先進的模型方面證明了它們的價值。你可以通過直接查看GPU 實現和CPU 實現來找到有關實現的更多詳細信息。
原文鏈接：
https://blog.tensorflow.org/2020/10/optimizing-tensorflow-lite-runtime.html