PyTorch - 卷积网络中的特征提取

PyTorch - 使用卷积神经网络 (ConvNets) 进行特征提取

卷积神经网络（Convolutional Neural Networks，简称 ConvNets 或 CNNs）在图像识别等任务中表现出色，因为它们能自动从数据中学习分层特征。特征提取是指利用预训练或自定义 ConvNet 的一部分来获取这些学习到的特征，然后将这些特征用于其他任务，例如迁移学习、图像相似度或作为其他机器学习模型的输入。

我们将探讨如何使用 PyTorch 进行特征提取，主要侧重于利用 torchvision.models 中的预训练模型。

步骤 1：导入必要的库

我们需要 torch、用于神经网络模块的 torch.nn、用于预训练架构的 torchvision.models 以及用于图像预处理的 torchvision.transforms。

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image # 用于加载示例图像

设置计算设备：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

步骤 2：加载预训练模型

torchvision.models 提供了许多预训练模型。我们将以 ResNet-18 为例。weights 参数允许我们加载预训练权重。使用 SomeModel_Weights.DEFAULT 是获取最佳可用预训练权重的现代方式。

# 加载预训练的 ResNet-18 模型
# 对于旧版本的 torchvision，你可能需要使用 pretrained=True
# model = models.resnet18(pretrained=True)
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model = model.to(device)
model.eval() # 将模型设置为评估模式

print("Original ResNet-18 model structure (first few layers and classifier):")
# print(model) # This will print the entire model structure

设置 model.eval() 非常重要，因为它会禁用 Dropout 和批量归一化（Batch Normalization）等层的训练行为，确保特征提取的一致性。

步骤 3：准备输入图像

预训练模型期望输入图像以特定方式进行预处理（例如，调整大小、归一化）。weights 对象通常提供推荐的转换（transforms）。

# 获取 ResNet-18 模型推荐的转换
transforms_pipeline = models.ResNet18_Weights.DEFAULT.transforms()
print("\nImage transforms pipeline:")
print(transforms_pipeline)

# 创建一个虚拟图像或加载一个真实图像
# 对于真实图像:
# try:
#   img = Image.open("your_image.jpg").convert("RGB")
# except FileNotFoundError:
#   print("Dummy image will be used as 'your_image.jpg' was not found.")
#   img = Image.new('RGB', (224, 224), color = 'red') # Dummy image

# 为了演示，我们创建一个表示单张图像批次的随机张量
dummy_image_tensor = torch.randn(1, 3, 224, 224).to(device) # 批量大小=1, 3 个通道, 224x224 像素

输入张量的形状通常应为 (批量大小, 通道数, 高度, 宽度)。

方法 1：使用模型子模块（直接访问层）

如果你想获取紧挨着最终分类层之前的层的特征，通常可以通过替换模型的分类器部分来修改模型。

# ResNet-18 最后的_全连接层_名为 'fc'
# 我们可以移除它，或用_恒等层_替换它，以获取其前一层的特征。
feature_extractor_v1 = nn.Sequential(*list(model.children())[:-1]) # 移除最后一层 (fc)
feature_extractor_v1 = feature_extractor_v1.to(device)
feature_extractor_v1.eval()

with torch.no_grad(): # 禁用推理时的梯度计算
    features_v1 = feature_extractor_v1(dummy_image_tensor)

print(f"\nFeatures from v1 (before fc layer, after avgpool): {features_v1.shape}")
# Expected output shape: torch.Size([1, 512, 1, 1]) for ResNet-18 (after adaptive avg pooling)
# We can flatten this: features_v1.squeeze().shape would be torch.Size([512]) for a single image

方法 2：使用前向钩子 (Forward Hooks)

前向钩子 (Forward Hooks) 是一种灵活的方法，可以在不修改模型架构的情况下获取任何中间层的激活值。钩子是一个函数，注册到模块上后，会在前向传播期间执行。

extracted_features = {}

def get_features_hook(module_name):
    """返回一个存储模块输出的钩子函数。"""
    def hook(module, input, output):
        # 对于卷积层，输出是一个 4D 张量 (批量, 通道数, 高度, 宽度)
        # 对于线性层，输出通常是一个 2D 张量 (批量, 特征数)
        extracted_features[module_name] = output.detach() # 分离张量以阻止梯度跟踪
    return hook

# 示例：从 ResNet-18 的 'layer3' 块的输出中提取特征
# 你需要知道层的名称。可以通过打印模型架构来找到它。
# print(model)

# 对于 ResNet-18，'layer3' 是一个 Sequential 块。我们可能想要获取它最后一个子块之后的特征。
# 我们以 model.layer3[-1] 为目标（layer3 中的最后一个 Bottleneck/BasicBlock）
target_layer_name = 'layer3_last_block'
handle = model.layer3[-1].register_forward_hook(get_features_hook(target_layer_name))

# 执行前向传播以触发钩子
with torch.no_grad():
    _ = model(dummy_image_tensor)

# 如果不再需要，在使用后移除钩子
handle.remove()

if target_layer_name in extracted_features:
    features_v2 = extracted_features[target_layer_name]
    print(f"\nFeatures from v2 (hook on '{target_layer_name}'): {features_v2.shape}")
else:
    print(f"\nCould not extract features for '{target_layer_name}'. Check layer name and model structure.")

理解代码

• model.eval()：对于确保 BatchNorm 和 Dropout 等层在推理期间表现正确至关重要。
• nn.Sequential(*list(model.children())[:-1])：这会创建一个新模型，包含原模型除最后一个子模块外的所有子模块。这是获取倒数第二层特征的常见方法。
• 前向钩子 (register_forward_hook)：一个函数注册到特定的模块（层）上。当前向传播时数据经过此模块，会调用钩子函数，并将模块自身、其输入和输出作为参数。这允许捕获中间激活值。
• output.detach()：在钩子中使用，用于获取一个不跟踪梯度的张量，这在不进行训练只做特征提取时非常重要。
• torch.no_grad()：一个上下文管理器，用于禁用梯度计算，从而减少内存消耗并加快推理速度。

实际应用

提取的特征可用于：

• 迁移学习：将这些特征用作新任务中一个较小的新模型的输入，特别是在新任务的数据有限时。
• 图像检索/相似度：比较不同图像的特征向量以查找相似图像。
• 可视化：t-SNE 或 UMAP 等技术可应用于特征向量，以可视化模型如何对不同类型的图像进行聚类。

选择从哪个层提取特征取决于具体任务。较早的层捕获更通用、低级的特征（边缘、纹理），而较深的层捕获更抽象、更高级的特征，这些特征特定于模型训练时的对象。