如何从 PyTorch 模型中获取特定层的输出?

问题描述

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

How to extract the features from a specific layer from a pre-trained PyTorch model (such as ResNet or VGG), without doing a forward pass again?

推荐答案

您可以注册一个 forward hook 在您想要的特定层上.类似的东西:

You can register a forward hook on the specific layer you want. Something like:

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):

For example, to obtain res5c output in ResNet, you may want to use a nonlocal variable (or global in Python 2):

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`.

这篇关于如何从 PyTorch 模型中获取特定层的输出?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！