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

Question

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

Answer 1

你可以在你想要的特定层上注册一个前向钩子。就像是：

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变量（或global在 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`.