如何将ViTImageProcessor嵌入Keras Sequential模型并解决训练报错

问题根因 #

你遇到的AttributeError: 'Tensor' object has no attribute 'ndim'，本质是自定义NormalizationLayer中直接调用ViTImageProcessor的预处理方法导致的。ViTImageProcessor的默认实现是针对numpy数组/PIL图像设计的，而Keras的.fit方法在图模式下会传入Tensor对象，处理器内部的逻辑（比如检查输入维度）会误将Tensor当作numpy数组处理，从而触发错误。

手动训练循环能运行，是因为你大概率在Eager模式下传入了numpy数组，避开了图模式下的Tensor兼容性问题，但代价是无法利用分布式策略和Keras的训练优化。

可行解决方案 #

1. 用Keras原生层替代ViTImageProcessor的归一化逻辑 #

直接从ViTImageProcessor中提取均值、标准差参数，用Keras官方的Normalization层实现归一化，完全兼容TensorFlow图模式：

from transformers import TFViTForImageClassification, ViTImageProcessor
import tensorflow as tf

# 自定义图像增强层（确保所有操作用TensorFlow API）
class AugmentationLayer(tf.keras.layers.Layer):
    def call(self, inputs, training=False):
        if not training:
            return inputs
        # 示例增强操作，可按需修改
        x = tf.image.random_flip_left_right(inputs)
        x = tf.image.random_brightness(x, max_delta=0.1)
        x = tf.image.random_contrast(x, lower=0.9, upper=1.1)
        return x

# 加载预训练处理器和ViT模型
processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
vit_model = TFViTForImageClassification.from_pretrained(
    "google/vit-base-patch16-224", num_labels=你的类别数
)

# 构建兼容图模式的Sequential模型
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(224, 224, 3)),  # 对应ViT的输入尺寸
    AugmentationLayer(),
    # 用Keras原生归一化层替代ViTImageProcessor的归一化
    tf.keras.layers.Normalization(
        mean=processor.image_mean,
        variance=[s**2 for s in processor.image_std]  # Normalization层需要方差而非标准差
    ),
    vit_model
])

# 编译模型
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=["accuracy"]
)

# 正常训练
model.fit(train_dataset, epochs=5, validation_data=val_dataset)

2. 若需完整的ViT预处理逻辑（如resize、裁剪） #

如果你的输入图像尺寸不固定，需要ViTImageProcessor的resize、中心裁剪等操作，建议将预处理逻辑整合到数据管道中，而非嵌入模型：

def preprocess_fn(image, label):
    # 将Tensor转为numpy数组，应用ViT处理器
    image_np = image.numpy()
    processed_image = processor(images=image_np, return_tensors="tf")["pixel_values"]
    return tf.squeeze(processed_image), label

# 增强逻辑也放到数据管道
def augment_fn(image, label):
    image = tf.image.random_flip_left_right(image)
    image = tf.image.random_brightness(image, max_delta=0.1)
    return image, label

# 构建数据管道
train_dataset = train_dataset.map(
    lambda x, y: tf.py_function(preprocess_fn, [x, y], [tf.float32, tf.int32]),
    num_parallel_calls=tf.data.AUTOTUNE
).map(augment_fn, num_parallel_calls=tf.data.AUTOTUNE).batch(32).prefetch(tf.data.AUTOTUNE)

# 此时模型只需加载ViT即可，无需嵌入预处理层
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(224, 224, 3)),
    vit_model
])

model.compile(...)
model.fit(train_dataset, ...)

3. 关键注意事项 #

• 所有自定义层必须继承tf.keras.layers.Layer，且内部操作仅使用TensorFlow API，禁止混用numpy操作（否则图模式下会报错）。
• 测试模型时，务必用Tensor输入（如tf.random.normal((1,224,224,3))）验证，而非仅用numpy数组，确保兼容图模式。
• 若要使用分布式策略，必须保证整个模型的所有层都是TensorFlow图兼容的，因此优先选择方案1的Keras原生层实现。

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