如何在PyTorch中生成VGG16第三、第四MaxPool层的中间输出

方法一：修改模型结构，直接返回指定层输出 #

这种方式适合需要固定提取这些层输出的场景，直接拆分原VGG16的特征提取模块，截取到目标MaxPool层即可。

先导入依赖并加载预训练模型：

import torch
import torchvision.models as models
from torchvision.transforms import Compose, ToTensor, Normalize

# 加载预训练VGG16模型
pretrained_model = models.vgg16(pretrained=True)
pretrained_model.eval()  # 切换到评估模式，避免BatchNorm/Dropout影响输出

定义自定义模型，截取到目标MaxPool层：

class CustomVGG(torch.nn.Module):
    def __init__(self, original_model):
        super().__init__()
        # 截取到第三层MaxPool（对应原features模块的前17层，索引0-16）
        self.features_to_pool3 = torch.nn.Sequential(*list(original_model.features.children())[:17])
        # 从第三层MaxPool后到第四层MaxPool的部分（索引17-23）
        self.features_to_pool4 = torch.nn.Sequential(*list(original_model.features.children())[17:24])
        
    def forward(self, x):
        pool3_out = self.features_to_pool3(x)
        pool4_out = self.features_to_pool4(pool3_out)
        return pool3_out, pool4_out

# 实例化自定义模型
custom_vgg = CustomVGG(pretrained_model)

准备输入并获取输出：

# 匹配VGG16训练时的图像预处理
transform = Compose([
    ToTensor(),
    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 模拟单张224x224的3通道输入（实际可替换为处理后的真实图像）
fake_img = torch.randn(1, 3, 224, 224)
processed_img = transform(fake_img) if isinstance(fake_img, torch.Tensor) else transform(fake_img).unsqueeze(0)

# 推理获取输出
with torch.no_grad():  # 评估阶段无需计算梯度
    pool3_output, pool4_output = custom_vgg(processed_img)

print("第三层MaxPool输出形状:", pool3_output.shape)
print("第四层MaxPool输出形状:", pool4_output.shape)

方法二：用前向钩子（Forward Hook）捕获输出 #

如果不想改动原模型结构，可借助PyTorch的钩子机制，在模型前向传播时自动记录指定层的输出。

import torch
import torchvision.models as models

# 加载预训练VGG16并设为评估模式
pretrained_model = models.vgg16(pretrained=True)
pretrained_model.eval()

# 存储中间输出的字典
outputs = {}

# 定义钩子函数：将层的输出存入字典
def hook_fn(module, input, output):
    outputs[module.__name__] = output

# 找到第三、四层MaxPool层并注册钩子（对应features模块的索引16和23）
pretrained_model.features[16].register_forward_hook(hook_fn)
pretrained_model.features[23].register_forward_hook(hook_fn)

# 准备输入并触发前向传播
fake_img = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    _ = pretrained_model(fake_img)  # 前向传播时钩子会自动记录输出

# 提取目标输出
pool3_output = outputs['16']
pool4_output = outputs['23']

print("第三层MaxPool输出形状:", pool3_output.shape)
print("第四层MaxPool输出形状:", pool4_output.shape)

注意事项 #

• 必须将模型切换到eval()模式，否则BatchNorm、Dropout等层会干扰输出结果。
• 真实图像输入时，必须严格按照VGG16的预处理规则做归一化，否则输出会偏离预期。
• 如果需要在训练过程中获取中间输出，去掉torch.no_grad()即可，钩子不会影响梯度计算。

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