Model Agnostic Approach to create a 2nd output head?

1

I’ve run into a strange problem where I need to take an arbitrary existing pytorch model and add a 2nd output head of the same kind as the original output head then have the model.forward() produce outputs from both output heads (e.g. head1_output, head2_output = model.forward(X)).
Alternatively it would also work if there was a way to increase the output channel/feature dimension of an existing arbitrary model.

Model Agnostic Constraints:

  • We cannot assume the outputs have a certain size which is consistent across inputs.
  • We also can’t assume particular model architecture (e.g. transformer vs CNN, vs fully-connected model).

Any ideas?

3

Hey thanks, the tricky part is that this might not work features = self.original_model(x). The original model probably evaluates the entire model back-bone and head and produces the output of the original head rather than a hidden representation.

So if we clone the original head it might not be compatible with the model’s outputs (e.g. Conv layer that has 64 input channels but 1 output channel cannot be re-applied to its own output).

4

Ignore the previous response as it’s a chatbot generated one and usually quite unhelpful.

1 Like