问题说明

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

相关阅读

评论可见,请评论后查看内容,谢谢!!!
 您阅读本篇文章共花了: