Record more in Recorder

fastai fastai dev
1

Hi Everyone,

I want to record the output activations at each layer in Recorder.

I have such a code block to register the forward hook:

outputs = {}

def save_output(name):
    def hook(module, input, output):
        outputs[name] = output
    return hook

for name, module in model.named_modules():
    if list(module.children()) == []:
        module.register_forward_hook(save_output(name))

And I could record the activations in pytorch ignite by:

@trainer.on(Events.ITERATION_COMPLETED)
def log_activations(engine):
    for key, value in outputs.items():
        print(key, value.data.mean().cpu().item(), value.data.std().cpu().item())

But, I don’t know how to implement it in fastai_v1 using callbacks, I have some questions:

  1. how to deal with the global variable outputs?
  2. where to put the code block that register the forward hooks?

Could you please help me? Thank you so much!

2

Callbacks are great for this - without messy globals and also handling removing the hooks automatically. I’ve just added some little classes to notebook 005 so you can simply do this:

image
image.png1206×1730 277 KB

This is HookCallback BTW:

class HookCallback(LearnerCallback):
    def on_train_begin(self, **kwargs):
        self.hooks = []
        for name,module in learn.model.named_modules():
            if list(module.children()) == []:
                func = self.hook(name,module)
                self.hooks.append(module.register_forward_hook(func))

    def on_train_end(self, **kwargs):
        if not self.hooks: return
        for hook in self.hooks: hook.remove()
        self.hooks=[]
3 Likes