PyTorch中FCN-ResNet50的forward hook多次触发问题求助

为什么hook触发三次？ #

你注册hook的model.backbone.layer4[2].relu是ResNet Bottleneck模块内的同一个ReLU实例，在Bottleneck的前向传播流程中，这个ReLU会被连续调用三次：

1. 第一次：conv1 -> bn1之后激活，输出shape [1,512,60,80]
2. 第二次：conv2 -> bn2之后激活，输出shape [1,512,60,80]
3. 第三次：残差相加之后激活，输出shape [1,2048,60,80]

每次调用ReLU都会触发你注册的forward hook，所以会打印三行输出。

解决办法 #

根据你的需求选择对应方案：

方案1：直接取最终需要的特征

如果你的目标是提取该Bottleneck模块的最终输出（也就是第三次ReLU的结果），不需要修改任何代码——你的my_output变量会被最后一次hook的结果覆盖，返回的就是shape [1,2048,60,80]的特征向量，这通常是我们需要的block级特征。

方案2：换hook目标，避免重复触发

直接把hook注册在Bottleneck模块本身，这样只会触发一次（对应整个模块的前向输出）：

_feature_vector = get_features_vector(path_to_image, model, model.backbone.layer4[2])

此时hook的output_就是该block的最终输出，和第三次ReLU的结果完全一致，只会打印一行输出。

方案3：精准获取某一次ReLU的输出

如果需要提取前两次ReLU的输出，可以在hook中加入计数逻辑，指定保留第几次的结果：

def get_features_vector(_path_img, _model, _layer, target_call=3) -> torch.Tensor:
    """
    Input:
        path_img: string, /path/to/image
        _model: a pretrained torch model
        target_call: int, 取第几次ReLU调用的输出
    Output:
        my_output: torch.tensor, 指定的输出
    """
    from PIL import Image
    import torchvision.transforms as transforms

    input_image = Image.open(_path_img)
    preprocess = transforms.Compose([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
        call_count = 0

        def my_hook(module_, input_, output_):
            nonlocal my_output, call_count
            call_count += 1
            print(f"Call {call_count}, Output Shape: {output_.shape}")
            if call_count == target_call:
                my_output = output_

        a_hook = _layer.register_forward_hook(my_hook)

        _model(input_batch)
        a_hook.remove()
        return my_output

调用时指定target_call参数即可，比如取第一次的输出：

_feature_vector = get_features_vector(path_to_image, model, model.backbone.layer4[2].relu, target_call=1)

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