如何在YOLOv7推理阶段提取并可视化任意层特征值？

一、修改Model类添加特征钩子逻辑 #

在models/yolo.py的Model类中添加特征存储和钩子注册功能，实现任意层特征的捕获：

1. 在Model类的__init__方法末尾添加属性初始化：

self.features = {}  # 用于存储提取到的特征，键为层索引字符串
self.target_layer_indices = []  # 记录需要提取的目标层索引

2. 给Model类新增一个方法，用于设置目标层并注册前向钩子：

def set_target_layers(self, indices):
    self.target_layer_indices = indices
    # 清除已有钩子，避免重复注册
    for module in self.module_list:
        if hasattr(module, '_forward_hooks'):
            module._forward_hooks.clear()
    
    # 定义钩子函数：捕获模块输出特征
    def hook_fn(module, input, output):
        idx = self.module_list.index(module)
        self.features[str(idx)] = output.detach().cpu()
    
    # 为目标层注册钩子
    for idx in indices:
        if 0 <= idx < len(self.module_list):
            self.module_list[idx].register_forward_hook(hook_fn)
        else:
            print(f"Warning: 层索引 {idx} 超出范围，当前模型共有 {len(self.module_list)} 层")

二、加载模型并指定目标层 #

在推理代码中加载模型，调用上述方法设置需要提取的层（比如你需要的102、103、104层）：

from models.yolo import Model
import torch

# 配置与权重路径
cfg_path = 'cfg/training/yolov7.yaml'
weights_path = 'yolov7.pt'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 初始化并加载模型
model = Model(cfg_path).to(device)
model.load_state_dict(torch.load(weights_path, map_location=device)['model'])
model.eval()

# 设置要提取的目标层索引（对应yaml中的层序号）
model.set_target_layers([102, 103, 104])

三、推理并提取特征 #

执行推理后，直接从model.features字典中获取对应层的特征：

# 预处理输入图像（替换为你自己的图像预处理逻辑）
# 示例：生成一个随机张量模拟输入，实际需用PIL/OpenCV读取图像后转成张量
input_img = torch.randn(1, 3, 640, 640).to(device)

# 推理（钩子会自动保存特征）
with torch.no_grad():
    _ = model(input_img)

# 获取目标层特征
feat_102 = model.features['102']
feat_103 = model.features['103']
feat_104 = model.features['104']

# 打印特征形状验证
print(f"第102层特征形状: {feat_102.shape}")
print(f"第103层特征形状: {feat_103.shape}")
print(f"第104层特征形状: {feat_104.shape}")

四、特征可视化实现 #

编写可视化函数，将特征图转换为可查看的图像（以灰度图为例）：

import matplotlib.pyplot as plt
import numpy as np

def visualize_feature_map(feature, layer_idx, show_channels=5):
    # 去除batch维度，得到(通道数, 高度, 宽度)的特征图
    feature_np = feature.squeeze(0).numpy()
    channel_num, h, w = feature_np.shape
    
    # 创建画布展示前N个通道的特征
    plt.figure(figsize=(16, 4))
    for i in range(min(show_channels, channel_num)):
        plt.subplot(1, show_channels, i+1)
        # 归一化到0-255范围
        norm_feat = (feature_np[i] - feature_np[i].min()) / (feature_np[i].max() - feature_np[i].min()) * 255
        plt.imshow(norm_feat.astype(np.uint8), cmap='gray')
        plt.title(f"Layer {layer_idx} - Channel {i}")
        plt.axis('off')
    plt.tight_layout()
    plt.show()

# 可视化第102层的前5个通道特征
visualize_feature_map(feat_102, 102)

注意事项 #

• 层索引对应关系：module_list的索引与你打印的模型结构中yaml层序号完全一致，确保你指定的索引正确。
• 多GPU适配：如果使用多GPU训练的权重，加载时需注意模型的module_list结构，必要时调整钩子注册逻辑。
• 显存优化：钩子中使用detach().cpu()将特征转移到CPU存储，避免占用GPU显存。

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