如何获取PyTorch版YOLOv5第一卷积层生成的特征图

以下是基于PyTorch框架提取YOLOv5第一卷积层特征图的具体步骤：

1. 加载YOLOv5模型与输入图像 #

首先导入所需依赖并加载预训练的YOLOv5模型，同时准备好输入图像：

from ultralytics import YOLO
from PIL import Image
import torchvision.transforms as transforms

# 加载YOLOv5s预训练模型（可替换为yolov5m/l/x等）
model = YOLO('yolov5s.pt')
# 加载测试图像
img = Image.open('test.jpg').convert('RGB')

# 图像预处理：匹配YOLOv5默认输入要求（640x640）
transform = transforms.Compose([
    transforms.Resize((640, 640)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
input_tensor = transform(img).unsqueeze(0)  # 添加batch维度

2. 定位第一卷积层 #

YOLOv5的模型结构存储在model.model列表中，第一个模块（model.model[0]）就是包含初始卷积层的Conv模块，其中的conv属性即为我们要捕获的第一卷积层：

# 获取第一卷积层
first_conv_layer = model.model[0].conv

3. 注册前向钩子捕获特征图 #

利用PyTorch的register_forward_hook机制，注册一个钩子函数来保存卷积层的输出特征图：

feature_map = None

def save_feature_map(module, input, output):
    global feature_map
    # 保存特征图并脱离计算图
    feature_map = output.detach()

# 给第一卷积层注册钩子
hook_handle = first_conv_layer.register_forward_hook(save_feature_map)

4. 执行前向传播获取特征图 #

将预处理后的图像输入模型，执行前向传播，钩子会自动捕获第一卷积层的输出：

# 前向传播（无需关心模型的检测输出）
model(input_tensor)

5. 可视化或保存特征图 #

此时feature_map已存储第一卷积层的输出，形状为(batch_size, channels, height, width)。可以提取单通道特征图进行可视化：

import matplotlib.pyplot as plt

# 取batch中第一个样本的第一个通道特征图
single_channel_feature = feature_map[0, 0, :, :].cpu().numpy()

# 可视化特征图
plt.imshow(single_channel_feature, cmap='gray')
plt.axis('off')
plt.show()

6. 清理钩子 #

使用完毕后移除钩子，避免内存泄漏：

hook_handle.remove()

注意事项 #

• 不同版本的YOLOv5（如旧版非ultralytics实现）可能模型结构略有差异，若model.model不存在，可尝试model.module或直接打印模型结构确认层级。
• 输入图像的尺寸需与模型默认输入尺寸一致（默认640x640），否则会导致特征图尺寸不符。

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