Add TimmWrapper
Author: qubvel
Pull request: yes
Labels: name: New model

name: Vision

name: run-slow

# What does this PR do?

Adds a TimmWrapper set of classes such that timm models can be loaded in as transformer models into the library. 

Continue of 
 - https://github.com/huggingface/transformers/pull/33687

### General Usage

```py
import torch
from urllib.request import urlopen
from PIL import Image
from transformers import AutoConfig, AutoModelForImageClassification, AutoImageProcessor

checkpoint = "timm/resnet50.a1_in1k"
img = Image.open(urlopen(
    'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
))

image_processor = AutoImageProcessor.from_pretrained(checkpoint)
inputs = image_processor(img, return_tensors="pt")
model = AutoModelForImageClassification.from_pretrained(checkpoint)

with torch.no_grad():
    logits = model(**inputs).logits

top5_probabilities, top5_class_indices = torch.topk(logits.softmax(dim=1) * 100, k=5)
```


### Pipeline

Timm models can now be used in the image classification (if a classification model) and image feature extraction pipelines

```py
import torch
from urllib.request import urlopen
from PIL import Image

from transformers import pipeline

img = Image.open(urlopen(
    'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
))
pipe = pipeline("image-classification", model="timm/resnet18.a1_in1k")
print(pipe(img))
```

### Trainer

Timm models can now be loaded and trained with the trainer class. 

Example model trained with the trainer running the script command below:
https://huggingface.co/qubvel-hf/vit-base-beans

```
python run_image_classification.py \                
    --dataset_name beans \
    --output_dir ./beans_outputs/ \
    --remove_unused_columns False \
    --label_column_name labels \
    --do_train \
    --do_eval \
    --push_to_hub \
    --push_to_hub_model_id vit-base-beans \
    --learning_rate 2e-5 \
    --num_train_epochs 5 \
    --per_device_train_batch_size 8 \
    --per_device_eval_batch_size 8 \
    --logging_strategy steps \
    --logging_steps 10 \
    --eval_strategy epoch \
    --save_strategy epoch \
    --load_best_model_at_end True \
    --save_total_limit 3 \
    --seed 1337 \
    --model_name_or_path timm/resnet18.a1_in1k \
    --ignore_mismatched_sizes
```

### Other features enabled

 - [x] Device map: 
 ```python
model = TimmWrapperForImageClassification.from_pretrained(checkpoint, device_map="auto")
 ```
 - [x] Torch dtype:
 ```python
model = TimmWrapperForImageClassification.from_pretrained(checkpoint, torch_dtype="bfloat16")
 ```
 - [x] Quantization:
 ```python
model = TimmWrapperForImageClassification.from_pretrained(checkpoint, load_in_4bit=True)
 ```
 - [x] Intermediate hidden states: `output_hidden_states=True` or `output_hidden_states=[1, 2, 3]` (to select specific hidden states)
```python
model = TimmWrapperForImageClassification.from_pretrained(checkpoint)
output = model(**intpus, output_hidden_states=True)
 ```
 
  - [x] Transformers TImmWrapper checkpoints are compatible with timm:
```python
model = timm.create_model("hf-hub:qubvel-hf/vit-base-beans", pretrained=True)
 ```
 
### TODO

 - [x] Gamma/beta renaming issue
 - [x] Update timm in CI 0.9.6 -> 1.0.11 to enable `output_hidden_states` tests
   - CI for slow-run takes longer to update images
 - [x] Weights are loaded by `transformers` instead of `timm`, which architectures are affected?
   - mostly OpenCLIP weights are affected (60-70 checkpoints), see the comment https://github.com/huggingface/transformers/pull/34564#issuecomment-2491972650
 - [x] Tests for image processor

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34564). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@rwightman could you please make the first review in case you have bandwidth

@qubvel I'm starting to work through it now, wanted to get eval working to check some familiar models and just wasted way too much time realizing I needed `--remove_unused_columns False ` for anything to work at all :/ ... that's a really poor setup when most datasets have an 'image' column and not a 'pixel_values' column (realize that's nothing to do with this PR, heh) :/

Annnyways, first pass of the code things looked sane but need to spend some time looking closer at the details and testing some cases.

A few high level q...

* Does 'Wrapper' add any worthwhile value/info in the name vs
  * `TimmPreTrainedModel(PreTrainedModel)`
  * `TimmModel(TimmPretrainedModel)`
  * `TimmModelForImageClassification(..)`

* Is there a reason TimmWrapperModel has `.timm_model`, instead of something more generic like `.model`

* Any reservations in changing TimmWrapperModelForImageClassification to not use TimmWrapperModel? There are a few issues with the handling details for classifier, possible optimizations for forward call sequence and it'd probably be a bit cleaner to just duplicate a bit of redunant code and keep the two impl separate and a bit different.

* I thought we were going to set pretained=True so timm can load the weights, their are a number of weight adaptation / translation things that don't run if this isn't use, cannot change num_classes cleanly for instance. 

* What happens if we try to push these models to the hub? Do they get uploaded/written in a form that timm can read?

@rwightman thanks for the review! Indeed there are some default params I'm also confused about, it's even more for object detection 🥲 

> Does 'Wrapper' add any worthwhile value/info in the name vs
TimmPreTrainedModel(PreTrainedModel)
TimmModel(TimmPretrainedModel)
TimmModelForImageClassification(..)

Left it as it was in the previous PR, however, `TimmModelForImageClassification` sounds better to me, I can rename it

> Is there a reason TimmWrapperModel has .timm_model, instead of something more generic like .model

The prefix "timm_model" is unique and is used in certain tests to identify when weights come from a timm model. It's also utilized in the `_fix_state_dict_key` method to determine whether to add the prefix when loading weights from the original checkpoint. For these reasons, I would prefer to keep it as "timm_model"

> Any reservations in changing TimmWrapperModelForImageClassification to not use TimmWrapperModel? There are a few issues with the handling details for classifier, possible optimizations for forward call sequence and it'd probably be a bit cleaner to just duplicate a bit of redunant code and keep the two impl separate and a bit different.

Originally, this was implemented without `TimmWrapperModel` in `TimmWrapperModelForImageClassification`, but I introduced it to reduce code repetition. This approach also aligns better with common patterns in the transformers repo. Could you provide more details on the issues with the classifier? In any case, the code will remain the same for both models if we aim to maintain `output_hidden_states` functionality across them.

> I thought we were going to set pretained=True so timm can load the weights, their are a number of weight adaptation / translation things that don't run if this isn't use, cannot change num_classes cleanly for instance.


I left a comment in the thread about this. We use transformers for weight loading to leverage features like device_map, torch_dtype, and quantization. I’m also unsure how to disable weight loading through transformers if it’s handled by timm, as I haven’t seen any examples of this in the repo. I can dig into it further if needed.

Do you have an estimate of how many models involve weight renaming? Is there a way to update checkpoints without breaking older versions of timm? Alternatively, could we manage similar weight renaming directly in transformers? (Though I think this approach may be less robust.)

> What happens if we try to push these models to the hub? Do they get uploaded/written in a form that timm can read?

I’m not sure if it's currently compatible, but it would be great to enable it! The config should be compatible, however, the weights state dict will have the "model.timm_model." prefix. I can look into removing this prefix before saving. I will try to enable it and add a separate test for a few checkpoints. Thanks for bringing this up!

> Originally, this was implemented without `TimmWrapperModel` in `TimmWrapperModelForImageClassification`, but I introduced it to reduce code repetition. This approach also aligns better with common patterns in the transformers repo. Could you provide more details on the issues with the classifier? In any case, the code will remain the same for both models if we aim to maintain `output_hidden_states` functionality across them.

I feel in this case there is a difference in alignment with other models, because both ImageClassification and base model wrap timm model instances that differ, instead of adding their own head to the same timm model. I see some options, depending on the mix of hidden state flags, head vs no head where it'd be appropriate to go through different forward calls, where more than just one argument might be appropriate to change on creation, etc.

Also, thinking about the future and other tasks. I feel there is a high probability say for supporting native timm object detection more flexibility is desired so it'd be safer to have it uncoupled.

That and there's a resisitance to significant changes in transformers after something is in there, so feel it's better to leave uncoupled to have additional flexibility and avoid being stuck with tricky decisions that might impact timm moving forward.


> > I thought we were going to set pretained=True so timm can load the weights, their are a number of weight adaptation / translation things that don't run if this isn't use, cannot change num_classes cleanly for instance.
> 
> I left a comment in the thread about this. We use transformers for weight loading to leverage features like device_map, torch_dtype, and quantization. I’m also unsure how to disable weight loading through transformers if it’s handled by timm, as I haven’t seen any examples of this in the repo. I can dig into it further if needed.
> 
> Do you have an estimate of how many models involve weight renaming? Is there a way to update checkpoints without breaking older versions of timm? Alternatively, could we manage similar weight renaming directly in transformers? (Though I think this approach may be less robust.)

Aside from renaming, it just doesn't work right now, you can't load weights for a model if the head size changes for a new classification task. The wrapper only works if you use the imagenet classifier.

I realize timm doesn't have the dtype, lazy init features but it's better to have it work I feel. Can potentially look at supporting some of that through timm. It hasn't been a priority as there aren't too many very large models in timm.

If there's no way to do use pretrained=True on creation then will probably need to figure out how to add & call a method in timm once transformers has the state_dict and before it's loaded into the model, not sure if there's a spot for such a call in transformers?

> Aside from renaming, it just doesn't work right now, you can't load weights for a model if the head size changes for a new classification task. The wrapper only works if you use the imagenet classifier.

Fixed!

> If there's no way to do use pretrained=True on creation then will probably need to figure out how to add & call a method in timm once transformers has the state_dict and before it's loaded into the model, not sure if there's a spot for such a call in transformers?

If there would be any timm model/class-specific function that fixes state_dict I suppose we can call it before loading weights, similar to `_fix_state_dict_key` in the current implementation. However, it will be supported only for newer timm versions

Thinking about this a bit more, there is another issue with the hub / transformers first weight loading. `timm` wasn't originally hub first, so the library itself is still the primary source of truth for some models, doing a hub based load won't work.

Example this model https://huggingface.co/laion/CLIP-ViT-B-16-laion2B-s34B-b88K ... is an OpenCLIP first model, but timm can load it if you use the model name 'vit_base_patch16_clip_224.laion2b' 

Indeed I feel that the wrapper should support all timm model names that work in timm, but right now if the model isn't on the hub w/ a timm primary config it isn't useable. Ideally it should work with both a hub model name OR any timm model name. The timm model name would require timm do the pretrained loading to resolve any translation to other hub name or weight source.

Some very popular examples of this

https://github.com/huggingface/pytorch-image-models/blob/51ac8d2efb926c6b7c34eeb1dc52bcf57999e2de/timm/models/vision_transformer.py#L1580-L1716

> > Aside from renaming, it just doesn't work right now, you can't load weights for a model if the head size changes for a new classification task. The wrapper only works if you use the imagenet classifier.
> 
> Fixed!

Yay, I was able to run fine-tune using run_image_classification.py after this fix. An observation the output files for that script aren't directly useable to push to hub.

There is no config matching timm format or name, there is a config output by the image preprocessor save process to a different filename that's a jumble of the timm config. 

Also, the state dict for the model has the timm_model prefix so it's not loadable in timm. Is there anyway to remove that prefix in the checkpoints? this would also make it more seamless doing local dir loads if someone had timm weights already and config files checked out. 

Hi @rwightman, 

I've addressed some comments and pushed the following fixes:  

Config and weights are now compatible with native timm models.  
  - Removed the `timm_model` prefix.  
  - Disabled processor config saving, the configuration is now kept in `config.json`, consistent with the original timm checkpoint format.  

Here’s an example of a fine-tuned **TimmWrapper model** using `run_classification` that can be loaded directly into timm:  

```python
model = timm.create_model("hf-hub:qubvel-hf/vit-base-beans", pretrained=True)
```  

Regarding OpenCLIP models:  
Would it be possible to provide compatible copies of these checkpoints under the `timm/` org on the Hub? Specifically, reconverting them to the appropriate format.  

Thanks!  

I've also noticed that some models from the provided [list](https://github.com/huggingface/pytorch-image-models/blob/51ac8d2efb926c6b7c34eeb1dc52bcf57999e2de/timm/models/vision_transformer.py#L1580-L1716) have transformers-like `config.json` along with `open_clip_config.json`, which allows them to be loaded with `CLIPForImageClassification`. Would it be possible to update other checkpoints similarly? We may confuse the users a bit, but updating model cards should help

- Saved checkpoints are now timm-compatible
- Removed inheritance from the `ForImageClassification` model and added additional tests
- The renaming of `TimmWrapperModel` to `TimmModel` is not yet done. I'm uncertain if this might cause issues, as the model folder would need to be named `timm`, which could potentially create confusion with the `timm` package, maybe it's safer to have it as timm_wrapper

@rwightman, could you please review when you have the bandwidth? also please let me know what you think about uploading transformers compatible checkpoints to new or existing repos on the hub (see comments above)

@qubvel thanks, I'll dig in and check out the latest updates tomorrow.

A user of the timm backbone was complaining about the transformers weight handling wrt to 'gamma', etc .. can the fixes for this use case be extended to the backbone easily?

> * Saved checkpoints are now timm-compatible
> * Removed inheritance from the `ForImageClassification` model and added additional tests
> * The renaming of `TimmWrapperModel` to `TimmModel` is not yet done. I'm uncertain if this might cause issues, as the model folder would need to be named `timm`, which could potentially create confusion with the `timm` package, maybe it's safer to have it as timm_wrapper
> 
> @rwightman, could you please review when you have the bandwidth? also please let me know what you think about uploading transformers compatible checkpoints to new or existing repos on the hub (see comments above)

Good point re the transformer naming resulting in a folder named just 'timm', yeah not a good idea, let's stick with wrapper naming as is. Parallels timm_backbone then too. 

> Hi @rwightman,
> 
> I've addressed some comments and pushed the following fixes:
> 
> Config and weights are now compatible with native timm models.
> 
> * Removed the `timm_model` prefix.
> * Disabled processor config saving, the configuration is now kept in `config.json`, consistent with the original timm checkpoint format.
> 
> Here’s an example of a fine-tuned **TimmWrapper model** using `run_classification` that can be loaded directly into timm:
> 
> ```python
> model = timm.create_model("hf-hub:qubvel-hf/vit-base-beans", pretrained=True)
> ```

Great, I will try this out.


> 
> Regarding OpenCLIP models: Would it be possible to provide compatible copies of these checkpoints under the `timm/` org on the Hub? Specifically, reconverting them to the appropriate format.
> 
> Thanks!

It's just OpenCLIP specific though most instances of this are OpenCLIP weights... there are a few other cases where an organization has uploaded a model and I've directly referened a checkpoint and leveraged conversion functions in timm. 

```
    "sam2_hiera_tiny.r896": _cfg(
        hf_hub_id='facebook/sam2-hiera-tiny',
        hf_hub_filename='sam2_hiera_tiny.pt',
    ),
```

Would be possible to do the conversions, there are 60-70ish cases of this I think so a bit of work.

> A user of the timm backbone was complaining about the transformers weight handling wrt to 'gamma', etc .. can the fixes for this use case be extended to the backbone easily?

yeah, probably the same trick will work, can have a look once this PR merged

@qubvel just noticed this reply, in review chat above, notifications in review chats always throw me off so I'm pulling it out here

> is the `last_hidden_state` would be equivalent to no pooled output?
> 
> where the last hidden state is:
> 
> ```python
> last_hidden_state, hidden_states = self.timm_model.forward_intermediates(pixel_values, **kwargs)
> # or
> last_hidden_state = self.timm_model.forward_features(pixel_values, **kwargs)
> ```
> 
> If so, I suppose we can add some flag to avoid applying `forward_head` if someone don't want to have pooler output. It can be the same`BaseModelOutputWithPooling` with `pooler_output=None`


Yes, last_hidden_state is always unpooled output, they are supposed to be equivalent btw `forward_intermediates` and `forward_features`,  last_hidden_state in both cases should be possible to pass to forward_head() and produce the same result. 

You've already noticed, but forward_features exists in every single model and forward_intermediates only a few of them.

So for pooling is `BaseModelOutputWithPooling` .. is it normal to have `pooler_output=None`? Looking again at some convnets and looks like they always do include pooling output, I assumed that some might output BaseModelOutput but seems that's only intermediate stages. So maybe this is fine as is if that's typical in transformers.

Which brings to next point, I'm not intimately famaliar with the norms in Transformers, we should get another reviewer on here who is as my focus is on the timm bits.


> So for pooling is BaseModelOutputWithPooling .. is it normal to have pooler_output=None? Looking again at some convnets and looks like they always do include pooling output, I assumed that some might output BaseModelOutput but seems that's only intermediate stages. So maybe this is fine as is if that's typical in transformers.

Agreed, looks a bit weird. In case of confusion, we can create `TimmWrapperModelOutput` and correctly specify that `pooler_output` is optional based on the config/init option. I will add an option and let`s see the core maintainers' opinion. 

Thanks, this looks good! cc @molbap can you give the processor code a quick look just to double check?

@LysandreJik @molbap @rwightman Thanks for the reviews! I believe all comments have been addressed. Do you have anything else in mind? It would be nice to move it forward.

Ok awesome! At this point just a second quick look from @ArthurZucker and we're good