前幾天看了TensorFlow官方微信推的 利用 XLA 將 GPU 性能推向極限 一文，里面提到使用Tensorflow xla可以顯著提升訓(xùn)練速度鄙币，在ResNet上速度最多可以提升3倍焰扳，而代碼的改動(dòng)是很少的辆琅。
XLA加速原理
按照google官方的說法测砂，xla的加速原理主要是融合內(nèi)核。看一個(gè)例子：
def model_fn（x环疼，y习霹，z）：
return tf.reduce_sum（x + y * z）
如果運(yùn)行模型時(shí)不使用 XLA，圖表會(huì)啟動(dòng)三個(gè)內(nèi)核炫隶，分別用于乘法淋叶、加法和減法。
使用了XLA后伪阶，它會(huì)將加法煞檩、乘法和減法 “融合” 到單個(gè) GPU 內(nèi)核中。這種融合運(yùn)算不會(huì)將 yz 和 x+yz 生成的中間值寫入內(nèi)存望门，而是將這些中間計(jì)算的結(jié)果直接 “流式傳輸” 給用戶形娇，并完整保存在 GPU 寄存器中。因?yàn)閯h除了內(nèi)存運(yùn)算筹误，所以能夠提升性能。
如何使用XLA
在官方提供的利用 XLA 將 GPU 性能推向極限一文中癣缅，提供了通過修改網(wǎng)絡(luò)運(yùn)算的方式來使用xla厨剪，這樣針對(duì)的是單個(gè)運(yùn)算，而且還必須先保證該運(yùn)算可以被xla編譯友存。這樣我覺得不太方便祷膳，有沒有可以自動(dòng)將網(wǎng)絡(luò)中所有可以被XLA編譯的運(yùn)算都改為使用XLA呢？翻了翻TensorFlow官方文檔屡立，的確是有的：
Turning on JIT compilation at the session level will result in all possible operators being greedily compiled into XLA computations. Each XLA computation will be compiled into one or more kernels for the underlying device. Subject to a few constraints, if there are two adjacent operators in the graph that both have XLA implementations, then they will be compiled into a single XLA computation. JIT compilation is turned on at the session level by setting the global_jit_level config to tf.OptimizerOptions.ON_1 and passing the config during session initialization.
里面提到了在session level配置jit可以把所有能被XLA編譯的運(yùn)算轉(zhuǎn)化為XLA運(yùn)算直晨。這樣就很方便了，只需要建立Session的時(shí)候配置好:
# Config to turn on JIT compilation
config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
sess = tf.Session(config=config)
這樣你就等于在使用xla加速了膨俐，十分方便勇皇！
