Forward hook is not called when using FX-based feature extractor

[yoshitomo-matsubara]

🐛 Describe the bug

This is an issue I talked to @datumbox

According to #4540, it looks like IntermediateLayerGetter will be replaced with FX-based feature extractor.
My ML OSS, torchdistill, is built on PyTorch / torchvision and heavily dependent on forward hook in PyTorch for knowledge distillation without modifying a model implementation to extract its intermediate feature.

However, the FX-based feature extractor disables some nn.Module features of its covered modules such as forward hook.
Here is a minimal example to demonstrate the issue.

import torch
from torch import nn
from torchvision import models
from torchvision.models.feature_extraction import create_feature_extractor
from torchvision.models._utils import IntermediateLayerGetter


def printnorm(self, input, output):
    # input is a tuple of packed inputs
    # output is a Tensor. output.data is the Tensor we are interested
    print('Inside ' + self.__class__.__name__ + ' forward')
    print('')
    print('input: ', type(input))
    print('input[0]: ', type(input[0]))
    print('output: ', type(output))
    print('')
    print('input size:', input[0].size())
    print('output size:', output.data.size())
    print('output norm:', output.data.norm())


resnet50 = models.resnet50()
return_layers = {'layer4': 'out'}
backbone_fx = create_feature_extractor(resnet50, return_layers)
backbone_org = IntermediateLayerGetter(resnet50, return_layers)


backbone_fx.layer1.register_forward_hook(printnorm)
backbone_org.layer1.register_forward_hook(printnorm)

x = torch.rand(2, 3, 224, 224)
# FX-based feature extractor doesn't call forward hook
z_fx = backbone_fx(x)

# IntermediateLayerGetter calls forward hook and print messages
z_org = backbone_org(x)
'''
Inside Sequential forward

input:  <class 'tuple'>
input[0]:  <class 'torch.Tensor'>
output:  <class 'torch.Tensor'>

input size: torch.Size([2, 64, 56, 56])
output size: torch.Size([2, 256, 56, 56])
output norm: tensor(1973.4232)
'''

Could you please enable the features for those wrapped by FX-based feature extractor?
Thank you!

Versions

Collecting environment information...
PyTorch version: 1.10.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 26 2021, 20:14:08) [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-94-generic-x86_64-with-glibc2.29
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-msssim==0.2.1
[pip3] torch==1.10.0
[pip3] torchdistill==0.1.4
[pip3] torchvision==0.11.1
[conda] Could not collect

[datumbox]

@yoshitomo-matsubara Thanks for raising this and for providing a snippet that show-cases the behaviour.

heavily dependent on forward hook in PyTorch for knowledge distillation without modifying a model implementation to extract its intermediate feature

I understand that previously you used register_forward_hook to extract the output of intermediate layers of models. In principle this can now be achieved using TorchVision's create_feature_extractor method. Is there another reason why you want to maintain the use of hooks at TorchDistill?

Could you please enable the features for those wrapped by FX-based feature extractor?

This behaviour originates from the use of FX tracing. As you can see by running the code below, the hook won't run on the traced model.

import torch
from torchvision import models

m = models.resnet50()
symbolic_traced = torch.fx.symbolic_trace(m)


def some_hook(self, input, output):
    print("Hook ran")
    raise Exception


symbolic_traced.layer1.register_forward_hook(some_hook)

x = torch.rand(2, 3, 224, 224)
z_fx = symbolic_traced(x)
print("Hook didn't run")

FX tracing flattens computations and replaces non built-in nn.Module layers, so this modifies the module of the model and makes some features not available. For example, you can no longer check modules with isinstance() because by default class info is not maintained. It's not clear to me whether hooks are supposed to play nice with FX or if they are not supported. Perhaps @jamesr66a could weight in here and let us know if that's something we currently support. As far as I can see on the documentation, there is no mention about the use of FX along with hooks.

[yoshitomo-matsubara]

Thank you @datumbox for the response.

I understand that previously you used register_forward_hook to extract the output of intermediate layers of models. In principle this can now be achieved using TorchVision's create_feature_extractor method. Is there another reason why you want to maintain the use of hooks at TorchDistill?

• Probably the biggest problem is that the use of FX changes the model interface (input-output) while we can use forward hook without any change in the interface For instance, when you train FX-wrapped models in end-to-end manner, you need to register the final layer as a key and extract the output from dict y_hat = model(x)['fc'] # if resnet whereas the classifier without FX-based extractor can do y_hat = model(x)
• (not limited to torchdistill) People use forward hook not only for feature extraction but also for exploring intermediate feature representations (input and output of specific layers) such as layer-wise latency analysis, dimensionality analysis and visualization without changing the interface.
• it looks like FX-based extractor captures only output of layers and misses input of layers

It would be very appreciated if you can support forward hook (and other features available for nn.Module if possible) even for FX-wrapped modules

[yoshitomo-matsubara]

Hi @jamesr66a
Could you please take a look at this issue?

[jamesr66a]

Hi @yoshitomo-matsubara,

torch.fx tracing only supports recording functional code that contributes data dependencies to the ultimate function outputs. print statements are side-effectful and are thus not recorded during tracing. If you'd like to preserve such side-effectful code, I recommend using torch.fx.wrap to create a function that wraps these operations and calling that within your forward hook, for example:

https://gist.github.com/jamesr66a/5d4e22e1008f76c8fc97d350b0dddceb

[yoshitomo-matsubara]

Thank you @jamesr66a for the clarification.

In your example, can't we do that by adding @torch.fx.wrap decorator to here function instead of having one more hop from the hook function to a function with the decorator (print_here function)
?

I wanted to extract features with dict in the hook function and added d = dict() to your example like this

import torch
import torch.fx

class Foo(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x)

class Bar(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.f = Foo()
    def forward(self, x):
        return self.f(x)

b = Bar()
d = dict()

@torch.fx.wrap
def print_here(output):
    print('here')
    d[2] = output
    return output


def here(self, input, output):
    d[1] = output
    return print_here(output)

b.f.register_forward_hook(here)

print('** run')
_ = b(torch.randn(5, 3))
print(d)
"""
** run
here
{1: tensor([[0.6329, 0.0000, 0.4976],
        [0.0000, 0.0000, 0.6403],
        [1.1902, 0.7864, 1.0433],
        [0.0000, 0.0000, 0.0000],
        [1.5052, 0.0000, 1.1598]]), 2: tensor([[0.6329, 0.0000, 0.4976],
        [0.0000, 0.0000, 0.6403],
        [1.1902, 0.7864, 1.0433],
        [0.0000, 0.0000, 0.0000],
        [1.5052, 0.0000, 1.1598]])}
"""

print('** trace')
traced = torch.fx.symbolic_trace(b)
print(d)
"""
** trace
{1: Proxy(relu), 2: tensor([[0.6329, 0.0000, 0.4976],
        [0.0000, 0.0000, 0.6403],
        [1.1902, 0.7864, 1.0433],
        [0.0000, 0.0000, 0.0000],
        [1.5052, 0.0000, 1.1598]])}
"""

print('** run trace')
_ = traced(torch.randn(5, 3))
print(d)
"""
** run trace
here
{1: Proxy(relu), 2: tensor([[0.0572, 0.4590, 0.0000],
        [0.0000, 0.0000, 1.2916],
        [0.0000, 0.0000, 1.0820],
        [0.0000, 0.0000, 0.0605],
        [0.9238, 0.1205, 0.1359]])}
"""

But it returns 1: Proxiy(relu) captured in the registered forward hook while a function one-hopped from the function (here) captures actual tensor with 2 as the key

I also tried the following here function with the decorator, but it still gives me 1: Proxy(relu) from the original hook function

@torch.fx.wrap
def here(self, input, output):
    d[1] = output
    return print_here(output)

[jamesr66a]

I think it's possible that you can @wrap the forward hook directly, but I used an indirection to be more safe

@yoshitomo-matsubara in general, it's pretty dangerous to be working with mutable, global data structures in conjunction with staging transforms (i.e. torch.fx). During symbolic tracing, the program is run with symbolic values (Proxy) solely for the purpose of capturing the program representation. Indeed, if you have side-effectful code (e.g. the dictionary mutation), that code will be triggered during tracing and may have unexpected behavior. Howver, if manipulation of the d dictionary happens purely inside the @wraped function, tracing should be safe since the body of a wrapped function will not be directly executed during tracing.

[yoshitomo-matsubara]

@jamesr66a Thank you for the advice.

I think it's possible that you can @Wrap the forward hook directly, but I used an indirection to be more safe

Actually, wrapping the forward hook function (i.e., here) with @torch.fx.wrap like the above example didn't work. The dictionary d still gets 1: Proxy(relu) )

If having the indirection in the original forward hook function resolves all the issues above, are there any side-effect or potential issues caused by wrapping the hook functions with @torch.fx.wrap registered on modules without torch.fx ?

Also a minor question: Your here function returns output, but the values returned from registered hook functions won't be used anywhere (or affect any subsequent process either), right?

[jamesr66a]

If having the indirection in the original forward hook function resolves all the issues above, are there any side-effect or potential issues caused by wrapping the hook functions with @torch.fx.wrap registered on modules without torch.fx ?

I don't believe so, @wrap simply registers the function in a dictionary for further modification during the tracing process: https://github.com/pytorch/pytorch/blob/8ba1ee6aa724f01736652c851e1a734e933f85ed/torch/fx/_symbolic_trace.py#L803. Outside of tracing, the original function is called directly and should be unaffected by the @wrap decorator

Also a minor question: Your here function returns output, but the values returned from registered hook functions won't be used anywhere (or affect any subsequent process either), right?

Right, they wont affect computation, but it's safer to have a data dependency from the output, since some transforms may erroneously remove the function when eliminating dead code (which is incorrect because the function is side-effectful). So returning the output value is defensive programming

[yoshitomo-matsubara]

Thank you @jamesr66a for the answers.

One last question/request; Can the tracing paradigm be updated to allow direct-wrap the forward hook? As mentioned above, it looks like the current implementation requires the indirection to get actual tensor instead of Proxy(...)