问题说明
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