794
學習筆記TF062:TensorFlow線性代數編譯框架XLA
XLA(Accelerated Linear Algebra)，線性代數領域專用編譯器(demain-specific compiler)，優化TensorFlow計算。即時(just-in-time,JIT)編譯或提前(ahead-of-time,AOT)編譯實現XLA，有助於硬件加速。XLA還在試驗階段。https://www.tensorflow.org/versions/master/experimental/xla/ 。
XLA優勢。線性代數領域專用編譯器，優化TensorFlow計算的執行速度(編譯子圖減少生命周期較短操作執行時間，融合管道化操作減少內存占用)、內存使用(分析、規劃內存使用需求，消除許多中間結果緩存)、自定義操作依賴(提高自動化融合底層操作low-level op性能，達到手動融合自定義操作custom op效果)、移動端內存占用(提前AOT編譯子圖減少TensorFlow執行時間，共享頭文件對被其他程序直接鏈接)、可移植性方麵(為新硬件開發新後端，TensorFlow不需要更改很多代碼用在新硬件設備上)。
XLA工作原理。LLVM編譯器框架係統，C++編寫，優化任意編程語言縮寫程序編譯時間(compile time)、鏈接時間(link time)、運行時間(run time)、空閑時間(idle time)。前端解析、驗證、論斷輸入代碼錯誤，解析代碼轉換LLVM中間表示(intermdediate representation,IR)。IR分析、優化改進代碼，發送到代碼生成器，產生本地機器代碼。三相設計LLVM實現。最重要，LLVM IR。編譯器IR表示代碼。C->Clang C/C++/ObjC前端、Fortran->llvm-gcc前端、Haskell->GHC前端 LLVM IR-> LLVM 優化器 ->LLVM IR LLVM X86後端->X86、LLVM PowerPC後端->PowerPC、LLVM ARM後端->ARM。https://www.aosabook.org/en/llvm.html 。
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計算。
config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
sess = tf.Session(config=config)
為一個或多個操作符手動打開JIT編譯。屬性_XlaCompile = true標記編譯操作符。
jit_scope = tf.contrib.compiler.jit.experimental_jit_scope
x = tf.placeholder(np.float32)
with jit_scope():
y = tf.add(x, x)
操作符放在XLA設備。有效設備XLA_CPU、XLA_GPU：
with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"):
output = tf.add(input1, input2)
JIT編譯MNIST實現。https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/mnist/mnist_softmax_xla.py 。
不使用XLA運行。
python mnist_softmax_xla.py --xla=false
運行完成生成時間線文件timeline.ctf.json，用Chrome跟蹤事件分析器 chrome://tracing，打開時間線文件，呈現時間線。左側列出GPU，可以看操作符時間消耗情況。
用XLA訓練模型。
TF_XLA_FLAGS=--xla_generate_hlo_graph=.* python mnist_softmax_xla.py
XLA框架處於試驗階段，AOT主要應用場景內存較小嵌入式設備、手機、樹莓派。
from future import absolute_import
from future import division
from future import print_function
import argparse
import sys
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.python.client import timeline
FLAGS = None
def main(_):
# Import data
mnist = input_data.read_data_sets(FLAGS.data_dir, one_hot=True)
# Create the model
x = tf.placeholder(tf.float32, [None, 784])
w = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
y = tf.matmul(x, w) + b
# Define loss and optimizer
y_ = tf.placeholder(tf.float32, [None, 10])
# The raw formulation of cross-entropy,
#
# tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(tf.nn.softmax(y)),
# reduction_indices=[1]))
#
# can be numerically unstable.
#
# So here we use tf.nn.softmax_cross_entropy_with_logits on the raw
# outputs of 'y', and then average across the batch.
cross_entropy = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y))
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
config = tf.ConfigProto()
jit_level = 0
if FLAGS.xla:
# Turns on XLA JIT compilation.
# 開啟XLA JIT編譯
jit_level = tf.OptimizerOptions.ON_1
config.graph_options.optimizer_options.global_jit_level = jit_level
run_metadata = tf.RunMetadata()
sess = tf.Session(config=config)
tf.global_variables_initializer().run(session=sess)
# Train
# 訓練
train_loops = 1000
for i in range(train_loops):
batch_xs, batch_ys = mnist.train.next_batch(100)
# Create a timeline for the last loop and export to json to view with
# chrome://tracing/.
# 在最後一次循環創建時間線文件，用chrome://tracing/打開分析
if i == train_loops - 1:
sess.run(train_step,
feed_dict={x: batch_xs,
y_: batch_ys},
options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE),
run_metadata=run_metadata)
trace = timeline.Timeline(step_stats=run_metadata.step_stats)
with open('timeline.ctf.json', 'w') as trace_file:
trace_file.write(trace.generate_chrome_trace_format())
else:
sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})
# Test trained model
correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
print(sess.run(accuracy,
feed_dict={x: mnist.test.images,
y_: mnist.test.labels}))
sess.close()
if name == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--data_dir',
type=str,
default='/tmp/tensorflow/mnist/input_data',
help='Directory for storing input data')
parser.add_argument(
'--xla', type=bool, default=True, help='Turn xla via JIT on')
FLAGS, unparsed = parser.parse_known_args()
tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
參考資料：
《TensorFlow技術解析與實戰》
歡迎推薦上海機器學習工作機會，我的微信：qingxingfengzi
最後更新：2017-11-11 01:03:36