且构网

分享程序员开发的那些事...
且构网 - 分享程序员编程开发的那些事

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

更新时间:2023-12-01 22:45:52

您可以注册一个 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`.