Register_forward_hook after every n steps

vision
1

I am training my own code, and I want to implement one function:

After training n(=100) iterations, I want to see feature maps.

Now I have done with register_forward_hook to see the feature maps, But I want to see them after n iterations. Any idea about how to implement this?

Thanks.

1 Like
2

I have implement it already…but a bit slow

3

You can keep track of the step or iteration by using a member variable inside the Network module or the module containing the hook function.
For example, inside your Network module,

# network module
class Network(nn.Module):
    def __init__(self):
        self.step = -1
# while training
model = Network()
for i in range(steps):
    # increment step var before forward pass
    model.step += 1
    output = model(input) # check self.step inside register_forward_hook
2 Likes
4

That’s an interesting approach!
Based on this, you could also create a “hook class” holding this internal parameter, which could be simpler than creating new modules:

class ActivationHook():
    def __init__(self, name, every_steps=1):
        self.name = name
        self.every_steps = every_steps
        self.steps = 0
        self.activation = {}
    
    def hook(self, model, input, output):
        if (self.steps % self.every_steps) == 0:
            self.activation[self.name] = output.detach()
        self.steps += 1
    

model = MyModel()
act_hook = ActivationHook('fc2', every_steps=2)
model.fc2.register_forward_hook(act_hook.hook)
x = torch.randn(1, 25)
output = model(x)
print(act_hook.activation['fc2'])
1 Like
5

This does not seem to work!