【问题标题】：How to get the output from a specific layer from a PyTorch model?如何从 PyTorch 模型中获取特定层的输出？
【发布时间】：2022-03-18 23:31:36
【问题描述】：

如何从预训练的 PyTorch 模型（例如 ResNet 或 VGG）中提取特定层的特征，而无需再次进行前向传递？

【问题讨论】：

标签： python pytorch

【解决方案1】：

新答案

编辑： there's a new feature in torchvision v0.11.0 that allows extracting features.

例如，如果你想从层 layer4.2.relu_2 中提取特征，你可以这样做：

import torch
from torchvision.models import resnet50
from torchvision.models.feature_extraction import create_feature_extractor

x = torch.rand(1, 3, 224, 224)

model = resnet50()

return_nodes = {
    "layer4.2.relu_2": "layer4"
}
model2 = create_feature_extractor(model, return_nodes=return_nodes)
intermediate_outputs = model2(x)

旧答案

您可以在您想要的特定层上注册forward hook。比如：

def some_specific_layer_hook(module, input_, output):
    pass  # the value is in 'output'

model.some_specific_layer.register_forward_hook(some_specific_layer_hook)
    
model(some_input)

例如，要在 ResNet 中获取 res5c 输出，您可能需要使用 nonlocal 变量（或 Python 2 中的 global）：

res5c_output = None

def res5c_hook(module, input_, output):
    nonlocal res5c_output
    res5c_output = output

resnet.layer4.register_forward_hook(res5c_hook)

resnet(some_input)
    
# Then, use `res5c_output`.

【讨论】：

这里输出的值是怎么返回的？？钩子函数不允许有返回值，所以我看不到代码中的 fc1000_output 将如何获得分配给它的输出值。 res5c_output 的值如何传递给 fc1000_output？
我添加了一个缺失的 nonlocal 声明。不是res5c_output的值被传递给fc1000_output，而是前一个变量绑定到了外部上下文。

【解决方案2】：

接受的答案非常有帮助！我在这里发布了一个完整的示例（使用@bryant1410 描述的注册钩子），供那些正在寻找工作解决方案的懒惰者使用：

import torch 
import torchvision.models as models
from torchvision import transforms
from PIL import Image

def get_feat_vector(path_img, model):
    '''
    Input: 
        path_img: string, /path/to/image
        model: a pretrained torch model
    Output:
        my_output: torch.tensor, output of avgpool layer
    '''
    input_image = Image.open(path_img)
    
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    
    input_tensor = preprocess(input_image)
    input_batch = input_tensor.unsqueeze(0)

    with torch.no_grad():
        my_output = None
        
        def my_hook(module_, input_, output_):
            nonlocal my_output
            my_output = output_

        a_hook = model.avgpool.register_forward_hook(my_hook)        
        model(input_batch)
        a_hook.remove()
        return my_output

你有你的特征提取函数，只需使用下面的sn-p调用它即可从resnet18.avgpool层获取特征

model = models.resnet18(pretrained=True)
model.eval()
path_ = '/path/to/image'
my_feature = get_feat_vector(path_, model)

【讨论】：