python - TensorFlow 无法使用 jit compile 运行（XLA 加速）

Question

平台：Windows 10
版本：TensorFlow-gpu 2.5、CUDA 11.2/11.1/11.0、cuDNN 8.1.0/8.1.1、Python 3.8
（我试过上面列出的不同版本的CUDA和cuDNN，每次安装新版本时，我总是卸载以前版本的所有组件，以防止由于多个版本而出现问题）

运行以下代码时：

import tensorflow as tf

print(tf.test.is_built_with_xla())
print(tf.test.is_built_with_gpu_support())
print(tf.config.list_physical_devices("GPU"))

我得到：

True
True
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

我也确信它在GPU加速下运行成功没有问题
（其实我从2.5发布就一直在使用这个，而在这个发布之前，我曾经使用2.1，并且运行没有问题）。

但是，例如，当我尝试使用XLA 加速时，代码如下：

import tensorflow as tf


@tf.function(jit_compile=True)
def operation(a, b):
    return a * b


print(operation(3, 4))

脚本无法运行并显示以下错误消息：

ptxas 在将 ptx 编译到 sass 期间返回错误：'Internal: ptxas exited with non-zero error code -1, output:' 如果错误消息表明无法写入文件，请验证是否提供了足够的文件系统空间。

我很确定有足够的空间（大约还剩 150GB，希望就足够了）

完整的日志消息如下：

53]成功打开动态库cudnn64_8.dll 2021-08-27 13:57:54.628215:I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871]添加可见gpu设备:0 2021-08-27 13:57: 54.628479：I tensorflow/core/platform/cpu_feature_guard.cc:142] 此 TensorFlow 二进制文件使用 oneAPI 深度神经网络库 (oneDNN) 进行了优化，可在性能关键操作中使用以下 CPU 指令：AVX AVX2 要在其他操作中启用它们，使用适当的编译器标志重建 TensorFlow。2021-08-27 13:57:54.628963: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] 找到具有属性的设备 0：pciBusID：0000:01:00.0 名称：NVIDIA GeForce GTX 1660 Ti 计算能力：7.5 coreClock ：1.59GHz coreCount：24 deviceMemorySize：6.00GiB deviceMemoryBandwidth：268.26GiB/s 2021-08-27 13:57:54.629231：169] XLA 服务 0x15d04eef0c0 为平台 CUDA 初始化（这不保证会使用 XLA）。设备：2021-08-27 13:57:55.174743：I tensorflow/compiler/xla/service/service.cc:177] StreamExecutor 设备 (0)：NVIDIA GeForce GTX 1660 Ti，计算能力 7.5 2021-08-27 13： 57：55.254238：F tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc:472] ptxas 在将 ptx 编译为 sass 期间返回错误：'内部：ptxas 以非零错误代码 -1 退出，输出：'如果错误消息表明无法写入文件，请验证是否提供了足够的文件系统空间。

进程以退出代码 -1073740791 (0xC0000409) 结束

标签： pythontensorflowdeep-learningtensorflow-xla