Dynamically add and remove layers through module hooks?

1

Hi

I’m trying to add adapters to ViT which I would love to be able to dynamically add and or remove.

I am currently able to add these layers to the graph through the hooks, and to freeze/unfreeze everything as needed, but using the handles to remove the hooks (and I expect the layers / parameters) doesn’t seem to work.

class Hooker:
    def __init__(self):
        self.hooks = {}
        #self.hk_idx = 0
    def register_hook(self, module, hook, layer_name):
        handle = module.register_forward_hook(hook)
        self.hooks[layer_name] = handle
        #self.hooks[self.hk_idx] = handle
        print(f"Hook registered for {layer_name}")
        print(f"Handle name: {handle}")
        self.hk_idx+=1

    def remove_hook(self, layer_name):
        if layer_name in self.hooks:
            self.hooks[layer_name].remove()
            print(f"Hook removed for {layer_name}")
            # del self.hooks[layer_name]
        else:
            print(f"No hook found for {layer_name}")

    def remove_all_hooks(self):
        for layer_name, handle in self.hooks.items():
            handle.remove()
            print(f"Hook removed for {layer_name}")
        # self.hooks = {}

When I count my parameters before and after adding the hooks, it correctly is adding the parameters, but when calling the function to remove them, they stay the same; am I missing something?

2

I don’t actually expect parameters to be added when hooks are registered to a module, do you know where that is happening?

3 
def add_adapters_with_hooks(image_encoder, bottleneck_dim):
    adapters = nn.ModuleList()
    hook_manager = Hooker()

    for idx, block in enumerate(image_encoder.blocks):
        adapter = Adapter(block.norm1.normalized_shape[0], bottleneck_dim)
        adapters.append(adapter)

        def hook(module, input, output, adapter=adapter):
            return adapter(output)

        hook_manager.register_hook(block, hook, f'block_{idx}')

    return adapters, hook_manager

and:

image_encoder = model.image_encoder
adapters, hook_manager = add_adapters_with_hooks(image_encoder, bottleneck_dim)
model.image_encoder.adapters = adapters
model.image_encoder.hook_manager = hook_manager
4

I see, thanks for clarifying. Since you manually appended the adapters into the module list, I’d expect you to also have to manually remove the adapters from the module list yourself. Removing hooks would not removing the parameters.

5

I’m not completely sure how I would about doing this, would you maybe know?

6

Maybe I am missing something, but it should be something like - as you append to the adapters, remember what index of the list you appended to. Then when you desire, you delete the corresponding adapter from the list.