如何提取预训练模型中第4和第10个Block层的输出？

方法一：修改模型forward函数 #

直接在模型的前向传播循环中，根据索引捕获目标Block的输出（注意：Python索引从0开始，第4个Block对应索引3，第10个对应索引9）：

def forward(self, x, ...):
    # 其他预处理代码...
    
    # 初始化字典存储目标输出
    block_outputs = {}
    
    # 遍历Block并记录指定层输出
    for idx, blk in enumerate(self.blocks):
        x = blk(x, B, T, W)
        # 捕获第4个Block输出
        if idx == 3:
            block_outputs["4th_block"] = x.detach()  # 若需要保留梯度可去掉detach()
        # 捕获第10个Block输出
        elif idx == 9:
            block_outputs["10th_block"] = x.detach()
    
    # 后续代码...
    return x, block_outputs  # 可同时返回最终输出和目标Block输出

调用模型时，就能同时拿到最终结果和两个指定Block的输出。

方法二：使用PyTorch钩子函数（无需修改模型结构） #

如果不想改动预训练模型的原代码，可通过注册前向钩子自动捕获目标Block的输出：

# 定义存储输出的字典
hook_outputs = {}

# 钩子函数：捕获模块输出并存储
def save_block_output(name):
    def hook(module, input, output):
        hook_outputs[name] = output.detach()
    return hook

# 为目标Block注册钩子（索引对应关系同上）
handle_4th = model.blocks[3].register_forward_hook(save_block_output("4th_block"))
handle_10th = model.blocks[9].register_forward_hook(save_block_output("10th_block"))

# 执行模型前向传播
final_output = model(your_input_data, B, T, W)  # 替换为你的输入参数

# 此时hook_outputs中已包含目标输出
print("第4个Block输出形状：", hook_outputs["4th_block"].shape)
print("第10个Block输出形状：", hook_outputs["10th_block"].shape)

# 用完钩子后移除，避免内存泄漏
handle_4th.remove()
handle_10th.remove()

这种方法无需修改原模型的forward逻辑，适合预训练模型的轻量化改造。

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