How can I get __add__ feature map?

#1

Hi, I have a model

class Model(nn.Module):
    def __init__(self):
        pass
    
    def forward(self, x):
        x1 = x * 2 
        x2 = x1 + x 
        x3 = x2 + 1
        return x3

I want to get x1, x2 and x3.
As a universal method, I CAN NOT return x1, x2, x3. I try forward_hook but it can only trace the module which defined in __init__.

#2

I have an idea.

  1. Find all add, sub, mul, div operators using fx
  2. Define Add, Sub, Mul, Div class.
  3. Add those modules using Module.add_module
  4. Run fx again and replace a+b to self.add(a, b)

Then those operators would be like a Module instead of a operator. So that I can use forward_hook.
BUT IT’S EXTREMELY COMPLEX， do you have any better suggestions? @ptrblck

#3

Can you be a bit more specific how your goal differs from what the fx show case for torchvision does?

The snippet I have in mind is:

from torchvision.models.feature_extraction import create_feature_extractor
# Confused about the node specification here?
# We are allowed to provide truncated node names, and `create_feature_extractor`
# will choose the last node with that prefix.
feature_extractor = create_feature_extractor(
	model, return_nodes=['blocks.0', 'blocks.1', 'blocks.2', 'blocks.3'])
# `out` will be a dict of Tensors, each representing a feature map
out = feature_extractor(torch.zeros(1, 3, 32, 32))

(And from the blog post, people do appear to believe that this is the best method to do it.)

Best regards

Thomas

#4

What a great new feature!
I woule like to have a try today, thank you so much. @tom

Edit: I have other more custom operators. create_feature_extractor is not enough but it inspires me. :smiley: Thanks again, tom.