TensorFlow与Keras报错：AttributeError: module 'tree' has no attribute 'flatten'

问题背景 #

安装最新版TensorFlow和Keras后，运行官方Fashion MNIST示例代码时触发AttributeError: module 'tree' has no attribute 'flatten'错误。

运行代码 #

# TensorFlow and tf.keras
import tensorflow as tf

# Helper libraries
import numpy as np
import matplotlib.pyplot as plt

fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

train_images, test_images = train_images/255, test_images/255

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])

model.compile(optimizer="adam",
              loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics = ["accuracy"])

model.fit(train_images, train_labels, epochs=10)

错误栈 #

Traceback (most recent call last):
  File "C:\Users\me\AppData\Local\Programs\Python\Python311\Tensorflow\example.py", line 21, in <module>
    tf.keras.layers.Input(shape=(28, 28)),
  File "C:\Users\me\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\src\layers\core\input_layer.py", line 143, in Input
    layer = InputLayer(
  File "C:\Users\me\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\src\layers\layer.py", line 216, in __new__
    obj = super().__new__(cls, *args, **kwargs)
  File "C:\Users\me\AppData\Local\Programs\Python\Python311\Lib\site-packages\keras\src\ops\operation.py", line 100, in __new__
    flat_arg_values = tree.flatten(kwargs)
AttributeError: module 'tree' has no attribute 'flatten'

已尝试无效的操作：

• 更换TensorFlow版本
• 升级tree模块
• 手动添加tree.flatten = lambda kwargs: kwargs

解决方案 #

1. 修复tree模块依赖冲突 #

问题核心是第三方可视化工具包tree与TensorFlow依赖的dm-tree（导入时用tree别名）冲突，执行以下命令：

# 卸载错误的tree包
pip uninstall -y tree
# 安装TensorFlow依赖的dm-tree
pip install dm-tree

2. 重新安装兼容的TensorFlow环境 #

单独安装Keras易导致版本不匹配，建议彻底清理后重新安装TensorFlow（自动包含兼容的Keras及所有依赖）：

# 卸载现有TensorFlow和Keras
pip uninstall -y tensorflow keras
# 重新安装最新版TensorFlow
pip install tensorflow
# 若Python 3.11存在兼容问题，指定安装支持3.11的版本（如2.13+）
# pip install tensorflow==2.13.0

3. 临时修改代码规避错误 #

若依赖修复暂时有困难，可修改模型定义，去掉Input层，直接在第一个Dense层指定input_shape：

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(28, 28)),
    tf.keras.layers.Dense(10)
])

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