使用TensorFlow目标检测API训练MobileNetV2触发JIT编译失败错误

• 运行环境：TensorFlow 2.9、Python 3.10、CUDA 11.7、cuDNN 8401
• 使用场景：基于TensorFlow Object Detection API训练自定义MobileNetV2模型，实现车辆图像的自定义目标检测功能
• 故障表现：执行训练启动命令后程序直接报错中断，无法进入正常训练流程

关键控制台输出如下：

W0615 19:12:26.293519 16220 deprecation.py:554] From C:\Users\oknor\AppData\Roaming\Python\Python310\site-packages\tensorflow\python\util\deprecation.py:629: calling map_fn_v2 (from tensorflow.python.ops.map_fn) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Use fn_output_signature instead
error: Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice
error: Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice
error: Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice
error: Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice
2022-06-15 19:13:01.007705: W tensorflow/core/framework/op_kernel.cc:1733] UNKNOWN: JIT compilation failed.
tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:
2 root error(s) found.
  (0) UNKNOWN:  JIT compilation failed.
         [[{{node train_input_images/write_summary/mod}}]]
         [[train_input_images/write_summary/Equal_1/_16]]
  (1) UNKNOWN:  JIT compilation failed.
         [[{{node train_input_images/write_summary/mod}}]]

1. 开头的map_fn_v2参数弃用警告属于TensorFlow版本迭代的正常提示，不会阻断训练运行，不是本次故障的诱因。
2. 核心根因是TensorFlow启动XLA JIT编译时，无法通过环境变量定位到CUDA路径下的nvvm/libdevice目录，找不到JIT编译依赖的libdevice库文件，直接触发编译失败，后续的图执行错误、节点报错都是该问题引发的连锁故障。

按优先级依次尝试以下方案：

• 先校验本地CUDA文件完整性：打开CUDA 11.7的安装目录（默认路径为C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7），检查nvvm\libdevice路径下是否存在libdevice.10.bc文件，如果文件缺失需要重新安装对应版本的CUDA。
• 配置系统环境变量：
  1. 新增系统变量CUDA_DIR，变量值填写CUDA 11.7的安装根路径，即上述默认路径中到v11.7的目录层级
  2. 编辑系统Path变量，新增%CUDA_DIR%\nvvm\libdevice条目，同时确认%CUDA_DIR%\bin、%CUDA_DIR%\libnvvp条目已存在
  3. 保存环境变量配置后，重启终端/IDE再重新执行训练命令
• 代码内强制指定路径（环境变量配置不生效时使用）：在训练脚本model_main_tf2.py的最开头，所有import语句之前添加如下代码，替换路径为本地实际CUDA安装路径：

import os
os.environ['CUDA_DIR'] = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7'
os.environ['PATH'] += r';C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\nvvm\libdevice'

• 临时规避方案：如果不需要XLA加速，直接在训练启动命令中追加参数--enable_xla=False关闭JIT编译，即可绕过该报错启动训练，仅会损失部分GPU训练速度。

内容的提问来源于stack exchange，提问作者ok Noronha