解決TF訓練提示 Not using XLA:CPU for cluster

訓練時一直未太關注該搞錯,啓動訓練後報警提示如下,瞭解了下XLA的設置,對性能有一定提升, 於是嘗試解決
  W tensorflow/compiler/jit/mark_for_compilation_pass.cc:1412] (One-time warning): Not using XLA:CPU for cluster because envvar TF_XLA_FLAGS=--tf_xla_cpu_global_jit was not set.  If you want XLA:CPU, either set that envvar, or use experimental_jit_scope to enable XLA:CPU.  To confirm that XLA is active, pass --vmodule=xla_compilation_cache=1 (as a proper command-line flag, not via TF_XLA_FLAGS) or set the envvar XLA_FLAGS=--xla_hlo_profile.

按照提示說的 要設置環境變量,參照ubuntu環境變量設置

vim  /etc/profile ,在末尾加上如下兩行

export XLA_FLAGS=--xla_hlo_profile
export TF_XLA_FLAGS=--tf_xla_cpu_global_jit

export PATH=$PATH:$XLA_FLAGS:$TF_XLA_FLAGS

source  /etc/profile

代碼中 參照TF 官方的 鏈接    

1、with tf.xla.experimental.jit_scope():
   c = tf.matmul(a, b) # compiled

2、with tf.xla.experimental.jit_scope(compile_ops=False): d = tf.matmul(a, c) # not compiled

3、with tf.xla.experimental.jit_scope( compile_ops=lambda node_def: 'matmul' in node_def.op.lower()): e = tf.matmul(a, b) + d # matmul is compiled, the addition is not.

https://tensorflow.google.cn/versions/r1.15/api_docs/python/tf/xla/experimental/jit_scope?hl=zh-cn#example_usage

參照第一條 設置後再次啓動訓練 

TF_XLA_FLAGS=--tf_xla_cpu_global_jit    python xxx.py   未再出現 Not using XLA:CPU提示

另外參照 一些文章裏提到 ,直接按如下 會報找不到 對應的flag ,可以設置到環境變量裏再試試

TF_XLA_FLAGS=--xla_generate_hlo_graph=.* python xxx.py
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章