Support multiple timm as backbones with forward_intermediates method

[ioangatop]

Feature request

Currently, in TimmBackbone requires that the timm model has a return_layers, which I'm not sure if any model implements:

transformers/src/transformers/models/timm_backbone/modeling_timm_backbone.py

Lines 80 to 83 in c15aad0

	# These are used to control the output of the model when called. If output_hidden_states is True, then
	# return_layers is modified to include all layers.
	self._return_layers = self._backbone.return_layers
	self._all_layers = {layer["module"]: str(i) for i, layer in enumerate(self._backbone.feature_info.info)}

Lately, timm implemented a forward_intermediates method to almost all networks for feature extraction - issue and pr.

Motivation

Thus, will little changes, we will be able to do

import transformers

model = transformers.Mask2FormerModel(
    config=transformers.Mask2FormerConfig(
        use_timm_backbone=True,
        backbone="vit_small_patch16_224",
    ),
)

as currently we get the error:

AttributeError: 'FeatureGetterNet' object has no attribute 'return_layers'

Your contribution

Review and/or implemented the PR

[amyeroberts]

Hi @ioangatop, thanks for raising this!

Sounds like a good idea to me. Would you like to open a PR to add these changes?

[ioangatop]

@amyeroberts thank you for the fast response

Will open a PR :)

[rwightman]

@ioangatop discussed this with @amyeroberts ... the quickest / safest approach right now would be to check the feature extractor wrapped model for return_layers with hasattr/getattr, set return_layers in the TimmBackbone module to None if it's not there.

In forward, if return_layers / all_layers is None should probably emit some sort of error if output_hidden_states is set to True, I'm not 100% if if the preferred approach to that is asserts or exceptions?

This would support the new FeatureGetter and also existing Feature wrappers (like Fx or Hooks) other than the default that also wouldn't work with this current approach. It's debatable that the output_hidden_states makes sense as implemented but that can be addresssed later, maybe providing better support with the new API for supported models...

[ioangatop]

@rwightman indeed, with a quick read (although I'm not that familiar with the transformers lib) the output_hidden_states should probably be refactored in a future PR (at least in TimmBackbone)...

I guess it has to be something like (?):

def __init__(...):
    ...
    # These are used to control the output of the model when called. If output_hidden_states is True, then
    # return_layers is modified to include all layers.
    self._backbone.return_layers = getattr(self._backbone, "return_layers", None)
    self._return_layers = self._backbone.return_layers
    self._all_layers = {layer["module"]: str(i) for i, layer in enumerate(self._backbone.feature_info.info)}
    ...

def forward(...):
    ...
    if output_hidden_states and self._return_layers is None:
        raise  ValueError("Cannot output hidden states for timm backbones")
    ...

Also for this, I guess we have to install timm from github main branch, at least until a new realise?

PS I haven't started the PR yet, but if someone has a more clear idea how to do it fast, I would be happy to just review it