Vision Transformer模型特征提取报错及模型摘要生成求助

问题描述 #

使用vit-pytorch实现的ViT模型编写图像特征提取代码时，遇到两个问题：

1. 执行hidden_layer = image_model(new_input)时触发错误：

AttributeError: EagerTensor object has no attribute 'reshape'.
If you are looking for numpy-related methods, please run the following:
from tensorflow.python.ops.numpy_ops import np_config
np_config.enable_numpy_behavior()

2. 无法正确构建可调用image_features_extract_model.summary()的特征提取模型

用户原代码：

image_model = ViT(
    image_size=224,
    patch_size=32,
    num_classes=1000,
    dim=1024,
    depth=6,
    heads=16,
    mlp_dim=2048,
    dropout=0.1,
    emb_dropout=0.1
)

# Image Shape is 224x224 RGB
input_shape = (1, 3, 224, 224)
output_shape = (1, 1024)

new_input = tf.keras.Input(shape=input_shape)
new_input = tf.squeeze(new_input, axis=0)
new_input = tf.zeros(input_shape)
new_input = tf.convert_to_tensor(new_input)
print(f'Input Shape : {new_input.shape}')
hidden_layer = image_model(new_input)
# 此处报错
print(f'\nOutput Shape : {hidden_layer.shape}')

image_features_extract_model = tf.keras.Model(new_input, output_shape, name="image_features_extract_model")

错误根源 #

1. 框架混用：vit-pytorch是PyTorch实现的模型，只能处理PyTorch张量，而代码中使用了TensorFlow生成的EagerTensor作为输入，模型内部调用PyTorch的reshape方法时，TensorFlow张量无此属性，导致报错。
2. Keras模型构建错误：tf.keras.Model的参数应为Input层和输出张量，而非输入张量和输出形状。

解决方案 #

方案一：全用PyTorch实现（推荐，适配原模型） #

直接使用PyTorch处理输入，构建特征提取模型，用torchinfo库实现类似Keras的summary功能：

import torch
from vit_pytorch import ViT
from torchinfo import summary

# 初始化ViT模型
image_model = ViT(
    image_size=224,
    patch_size=32,
    num_classes=1000,
    dim=1024,
    depth=6,
    heads=16,
    mlp_dim=2048,
    dropout=0.1,
    emb_dropout=0.1
)

# 定义特征提取模型：移除原分类头，保留[CLS] token的输出
class FeatureExtractor(torch.nn.Module):
    def __init__(self, vit_model):
        super().__init__()
        self.vit = vit_model
        # 将原分类头替换为恒等层，丢弃分类逻辑
        self.vit.head = torch.nn.Identity()
    
    def forward(self, x):
        return self.vit(x)

# 构建特征提取模型
feature_extractor = FeatureExtractor(image_model)

# 测试输入：PyTorch张量，形状为(batch_size, channels, height, width)
input_tensor = torch.randn(1, 3, 224, 224)
hidden_layer = feature_extractor(input_tensor)
print(f'Input Shape : {input_tensor.shape}')
print(f'Output Shape : {hidden_layer.shape}')

# 打印模型结构（需先安装torchinfo：pip install torchinfo）
summary(feature_extractor, input_size=(1, 3, 224, 224))

方案二：转用TensorFlow/Keras版ViT #

如果坚持使用TensorFlow生态，替换为TensorFlow实现的ViT模型，正常构建Keras特征提取模型：

import tensorflow as tf
from tensorflow.keras.applications import ViT

# 使用TensorFlow官方ViT模型（需先安装tensorflow-models：pip install tensorflow-models）
image_model = ViT(
    image_size=224,
    patch_size=32,
    num_classes=1000,
    projection_dim=1024,
    num_layers=6,
    num_heads=16,
    mlp_dim=2048,
    dropout_rate=0.1,
    representation_size=None  # 直接取[CLS] token输出
)

# 构建特征提取模型：输入层连接ViT，丢弃分类头
input_layer = tf.keras.Input(shape=(224, 224, 3))  # TensorFlow默认通道在后
# 关闭训练模式，避免更新模型参数
feature_output = image_model(input_layer, training=False)
feature_extractor = tf.keras.Model(input_layer, feature_output, name="image_features_extract_model")

# 测试输入
input_tensor = tf.random.normal((1, 224, 224, 3))
hidden_layer = feature_extractor(input_tensor)
print(f'Input Shape : {input_tensor.shape}')
print(f'Output Shape : {hidden_layer.shape}')

# 打印模型结构
feature_extractor.summary()

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