天天看點

學習筆記TF062:TensorFlow線性代數編譯架構XLA

XLA優勢。線性代數領域專用編譯器,優化TensorFlow計算的執行速度(編譯子圖減少生命周期較短操作執行時間,融合管道化操作減少記憶體占用)、記憶體使用(分析、規劃記憶體使用需求,消除許多中間結果緩存)、自定義操作依賴(提高自動化融合底層操作low-level op性能,達到手動融合自定義操作custom op效果)、移動端記憶體占用(提前AOT編譯子圖減少TensorFlow執行時間,共享頭檔案對被其他程式直接連結)、可移植性方面(為新硬體開發新後端,TensorFlow不需要更改很多代碼用在新硬體裝置上)。

XLA輸入語言HLO IR,XLA HLO定義圖形,編譯成各種體系結構機器指令。編譯過程。XLA HLO->目标無關優化分析->XLA HLO->XLA後端->目标相關優化分析->目标特定代碼生成。XLA首先進行目标無關優化分析(公共子表達式消除common subexpression elimination CSE,目标無關操作融合,運作時記憶體緩沖區分析)。XLA将HLO計算發送到後端。後端執行進一步HLO級目标不相關優化分析。XLA GPU後端執行對GPU程式設計模型有益操作融合,确定計算劃分成流。生成目标特定代碼。XLA CPU、GPU後端用LLVM中間表示、優化、代碼生成。後端用LLVM IR表示XLA HLO計算。XLA 支援x86-64、NVIDIA GPU JIT編譯,x86-64、ARM AOT編譯。AOT更适合移動、嵌入式深度學習應用。

JIT編譯方式。XLA編譯、運作TensorFlow計算圖一部分。XLA 将多個操作(核心)融合到少量編譯核心,融合操作符減少存儲器帶寬提高性能。XLA 運作TensorFlow計算方法。一,打開CPU、GPU裝置JIT編譯。二,操作符放在XLA_CPU、XLA_GPU裝置。

打開JIT編譯。在會話打開。把所有可能操作符程式設計成XLA計算。

為一個或多個操作符手動打開JIT編譯。屬性_XlaCompile = true标記編譯操作符。

操作符放在XLA裝置。有效裝置XLA_CPU、XLA_GPU:

不使用XLA運作。

運作完成生成時間線檔案timeline.ctf.json,用Chrome跟蹤事件分析器 chrome://tracing,打開時間線檔案,呈現時間線。左側列出GPU,可以看操作符時間消耗情況。

用XLA訓練模型。

XLA架構處于試驗階段,AOT主要應用場景記憶體較小嵌入式裝置、手機、樹莓派。

參考資料:

《TensorFlow技術解析與實戰》

歡迎推薦上海機器學習工作機會,我的微信:qingxingfengzi

繼續閱讀