TensorFlow 2.6 GPU版dtensor导入失败问题求助

环境信息 #

• Anaconda3 Spyder环境，Python 3.9
• TensorFlow 2.6（GPU版），已配置CUDA 11.2 + cuDNN 8，tensorflow.test.is_gpu_available()返回True

问题描述 #

导入tensorflow.keras.callbacks.TensorBoard时触发ImportError，报错提示无法从tensorflow.compat.v2.experimental导入dtensor。

报错详情 #

File "C:\Users\Tibame_EX14\.spyder-py3\GPU\UsingTensorflowBasic.py", line 18, in <module>
    from tensorflow.keras.callbacks import TensorBoard

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\api\_v2\keras\__init__.py", line 12, in <module>
    from keras import __version__

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\__init__.py", line 21, in <module>
    from keras import models

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\models\__init__.py", line 18, in <module>
    from keras.engine.functional import Functional

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\engine\functional.py", line 26, in <module>
    from keras import backend

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\backend.py", line 34, in <module>
    from keras.engine import keras_tensor

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\engine\keras_tensor.py", line 19, in <module>
    from keras.utils import object_identity

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\utils\__init__.py", line 48, in <module>
    from keras.utils.layer_utils import get_source_inputs

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\utils\layer_utils.py", line 26, in <module>
    from keras import initializers

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\initializers\__init__.py", line 22, in <module>
    from keras.initializers import initializers_v2

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\initializers\initializers_v2.py", line 23, in <module>
    from keras.dtensor import utils

  File "C:\Users\Tibame_EX14\anaconda3\lib\site-packages\keras\dtensor\__init__.py", line 22, in <module>
    from tensorflow.compat.v2.experimental import dtensor as dtensor_api

ImportError: cannot import name 'dtensor' from 'tensorflow.compat.v2.experimental' (C:\Users\Tibame_EX14\anaconda3\lib\site-packages\tensorflow\_api\v2\compat\v2\experimental\__init__.py)

解决步骤 #

1. 匹配Keras与TensorFlow版本 #

TensorFlow 2.6对应Keras 2.6.x版本，版本不匹配是最常见原因。执行以下命令修复：

# 查看当前Keras版本
pip show keras
# 卸载现有Keras
pip uninstall -y keras
# 安装对应版本
pip install keras==2.6.0

2. 重新安装TensorFlow #

如果TensorFlow的API文件损坏，重新安装即可：

# 卸载现有TensorFlow（GPU版用tensorflow-gpu）
pip uninstall -y tensorflow
# 安装指定版本
pip install tensorflow==2.6.0

GPU版替换为：

pip uninstall -y tensorflow-gpu
pip install tensorflow-gpu==2.6.0

3. 清理Anaconda缓存并重装 #

Anaconda缓存可能导致版本冲突，执行：

conda clean --all
conda uninstall tensorflow keras
conda install tensorflow-gpu==2.6.0 keras==2.6.0

4. 临时替代导入 #

紧急情况下可以用以下方式导入TensorBoard：

from tensorflow.python.keras.callbacks import TensorBoard

内容的提问来源于stack exchange，提问作者Shuhan You