Custom loss based on old model weights

Hi there,

This question concerns something I already have working, but a lot is dependent on how the model is defined. I would like some advice on what a better implementation would be such that it works with other (fully connected / convolutional) models and layers.

Currently I am working on a supplementary custom loss that is based on activation and weights. These weights and activation can be taken in various ways, e.g. the weights of a Linear layer and its direct activation output, the weights of a Conv2d layer and the feature maps thereafter, or even the patches of input to a single Conv2d filter and its weights. There are about 3 - 6 combinations possible for what I’m working on, let’s call them type A, type B, etc…

Because this custom loss requires activation from specific (parts of) layers/blocks in the model, what would be the best way to implement this such that it best generalizes to most or any (fully connected / convolutional) Pytorch model? The implementation I am working with uses hooks that read out activation and weights from parts of the network that matches the given type. The hook that takes activation from the first ReLU after a Conv2d layer might look something like this:

    def set_activation_hooks(self):
        conv_layer_id = None
        for name, layer in self.model.named_children():
            if isinstance(layer, Conv2d):
                conv_layer_id = name
            if isinstance(layer, ReLU) and conv_layer_id is not None:
                hook = layer.register_forward_hook(self.save_outputs_hook(conv_layer_id, name))
                self.activation_hooks.append(hook)
                conv_layer_id = None

However, this obviously breaks down when the model is structured even the tiniest bit differently: when for example the ReLU is part of the forward pass or when it’s part of a Sequential block. Are hooks (like this) really the way to go, or are there better alternatives?

Thanks in advance!