彻底解决conda环境下 tensorflow gpu版本出现的问题：Not creating XLA devices, tf_xla_enable_xla_devices not set

问题说明

tf.test.is_gpu_available()结果为TRUE，但有Not creating XLA devices, tf_xla_enable_xla_devices not set报错

XLA

XLA的全称是Accelerated Linear Algebra，即加速线性代数。作为一种深度学习编译器，长期以来被作为Tensorflow框架的一个试验特性被开发，历时至今已经超过两三年了，随着Tensorflow 2.X的发布，XLA也终于从试验特性变成了默认打开的特性。此外， Pytorch社区也在大力推动XLA在Pytorch下的开发，现在已经有推出PyTorch/XLA TPU版本，暂只支持谷歌平台TPU上使用。

分析原因

可见XLA是一个GPU加速模块，如果你的tf.test.is_gpu_available()结果为TRUE，TensorFlow也可使用GPU进行训练，但没打开XLA进行加速。2.4.0版本XLA功能默认关闭，官方更新文档给出了解决方案

即需要在系统环境变量中添加，使其默认打开XLA加速。博客https://blog.csdn.net/FriendshipTang/article/details/114685263给出了一种在代码前进行添加的方法

import tensorflow as tf
import os
os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
print(tf.__version__)
a = tf.constant(1.)
b = tf.constant(2.)
print(a+b)
print('GPU:', tf.test.is_gpu_available())

os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices'
类似于宏定义，但每次写代码都需要重新申明，有没有修改虚拟环境的系统变量的方法呢

最终解决

可以使用conda env config vars set TF_XLA_FLAGS=--tf_xla_enable_xla_devices设置当前虚拟环境中的自定义环境变量。设置完毕后需重新启动虚拟环境才可生效，再次使用tf.test.is_gpu_available()便不会再出现Not creating XLA devices, tf_xla_enable_xla_devices not set`报错，且训练速度会大大加快，可以使用以下代码进行速度测试。

import tensorflow as tf
import timeit

with tf.device('/cpu:0'):
    cpu_a = tf.random.normal([10000, 1000])
    cpu_b = tf.random.normal([1000, 2000])
    print(cpu_a.device, cpu_b.device)

with tf.device('/gpu:0'):
    gpu_a = tf.random.normal([10000, 1000])
    gpu_b = tf.random.normal([1000, 2000])
    print(gpu_a.device, gpu_b.device)

def cpu_run():
    with tf.device('/cpu:0'):
        c = tf.matmul(cpu_a, cpu_b)
    return c

def gpu_run():
    with tf.device('/gpu:0'):
        c = tf.matmul(gpu_a, gpu_b)
    return c

# warm up	这里就当是先给gpu热热身了
cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('warmup:', cpu_time, gpu_time)


cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('run time:', cpu_time, gpu_time)

参考博客

https://github.com/tensorflow/tensorflow/issues/44683
https://blog.csdn.net/weixin_42642296/article/details/112565119
https://blog.csdn.net/robot8me/article/details/109471568
https://blog.csdn.net/FriendshipTang/article/details/114685263