如何使CNN输出特征向量？基于Python实现类Flatten层的图像特征提取方法咨询

嘿，这事儿其实挺简单的——本质就是做特征提取，跳过CNN后面的分类全连接层就行。我给你用目前最常用的两个框架分别写示例，你照着改就行，毕竟你已经懂CNN结构了，一看就明白。

首先我们先搭一个基础的CNN结构，然后把分类头砍掉，只保留到Flatten（或者你也可以用全局平均池化GlobalAveragePooling2D，输出维度更低）的部分：

import tensorflow as tf
from tensorflow.keras import layers, models

# 1. 构建带分类头的完整CNN（方便你对比）
def build_full_cnn(input_shape=(224,224,3), num_classes=10):
    model = models.Sequential([
        layers.Conv2D(32, (3,3), activation='relu', input_shape=input_shape),
        layers.MaxPooling2D((2,2)),
        layers.Conv2D(64, (3,3), activation='relu'),
        layers.MaxPooling2D((2,2)),
        layers.Conv2D(128, (3,3), activation='relu'),
        layers.MaxPooling2D((2,2)),
        # 这里就是你要的类似Flatten的层
        layers.Flatten(),
        # 下面是分类头，我们后面会单独剥离特征提取部分
        layers.Dense(128, activation='relu'),
        layers.Dense(num_classes, activation='softmax')
    ])
    return model

# 2. 创建纯特征提取模型（只取到Flatten层的输出）
full_model = build_full_cnn()
# 你可以先跑full_model.summary()查看所有层的名字，替换成你想要的输出层
feature_extractor = models.Model(inputs=full_model.input, 
                                 outputs=full_model.get_layer('flatten').output)

# 3. 测试提取特征
# 实际使用时把随机张量换成你预处理后的图像（注意尺寸和input_shape匹配）
test_img = tf.random.normal((1, 224, 224, 3))  # batch_size=1的输入
feature_vector = feature_extractor.predict(test_img)

print(f"特征向量维度: {feature_vector.shape}")  # 输出维度取决于你的池化次数和卷积核数量

如果不想先搭完整模型，也可以直接用函数式API写纯特征提取模型，更灵活：

def build_feature_extractor(input_shape=(224,224,3)):
    inputs = layers.Input(shape=input_shape)
    x = layers.Conv2D(32, (3,3), activation='relu')(inputs)
    x = layers.MaxPooling2D((2,2))(x)
    x = layers.Conv2D(64, (3,3), activation='relu')(x)
    x = layers.MaxPooling2D((2,2))(x)
    x = layers.Conv2D(128, (3,3), activation='relu')(x)
    x = layers.MaxPooling2D((2,2))(x)
    outputs = layers.Flatten()(x)  # 这就是最终输出的特征向量
    model = models.Model(inputs=inputs, outputs=outputs)
    return model

# 直接使用这个特征提取模型
feature_extractor = build_feature_extractor()
feature_vector = feature_extractor(test_img)

PyTorch的思路类似，我们可以在模型的forward方法里加个开关，控制输出特征还是分类结果；也可以直接剥离特征提取的子模块：

import torch
import torch.nn as nn

# 1. 定义包含特征提取和分类头的CNN
class CNN(nn.Module):
    def __init__(self, input_channels=3, num_classes=10):
        super(CNN, self).__init__()
        # 特征提取部分
        self.features = nn.Sequential(
            nn.Conv2d(input_channels, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Flatten()  # 特征向量输出层
        )
        # 分类头部分
        self.classifier = nn.Sequential(
            nn.Linear(128 * 28 * 28, 128),  # 224经过3次池化后尺寸是28x28
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x, return_feature=False):
        feature = self.features(x)
        if return_feature:
            return feature  # 返回特征向量
        else:
            return self.classifier(feature)  # 返回分类结果

# 2. 初始化模型并提取特征
model = CNN()
test_img = torch.randn(1, 3, 224, 224)  # PyTorch输入格式是(batch_size, channels, h, w)

# 直接获取特征向量
feature_vector = model(test_img, return_feature=True)
print(f"特征向量维度: {feature_vector.shape}")

如果你只想单独用特征提取模块，也可以直接调用：

feature_extractor = model.features
feature_vector = feature_extractor(test_img)

小提示 #

• 要是用预训练模型（比如ResNet、VGG），思路完全一样：砍掉最后全连接层，提取中间层的输出就行
• 预处理时一定要注意图像尺寸、归一化方式和模型要求匹配（比如预训练模型的均值方差）
• 如果觉得Flatten后的维度太高，可以用GlobalAveragePooling2D（Keras）或者nn.AdaptiveAvgPool2d(1)（PyTorch）得到固定维度的特征向量，实用性更强

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