Visual Encoders (ViT, DeiT) Support Rollout - #1546
Merged
Merged
Conversation
Implement ViT/DeiT architecture adapter for model bridging.
Implement VisionClassifierHeadBridge to handle CLS token slicing for classification.
Updated documentation for pixel_values parameter to clarify its use with vision models.
Added vision model architectures and classification heads.
Updated the VisionClassifierHeadBridge to directly use an already-pooled CLS token instead of slicing from the sequence output. Adjusted the forward method to reflect this change and improved error handling for the original component.
This file contains unit tests for the ViTArchitectureAdapter, covering component mapping, configuration flags, weight conversions, and model preparation methods.
Reintroduce a patch_layers function to inject a dummy 'mlp' attribute into ViTLayer blocks for MLPBridge compatibility.
Removed the patch_layers function and its call, which injected a dummy 'mlp' attribute into ViTLayer blocks. Updated comments for clarity regarding the MLPBridge container.
Inject a dummy 'mlp' attribute into ViTLayer blocks to satisfy hasattr check for TransformerLens.
Contributor
Author
|
All checks pass now, ready to merge |
jlarson4
reviewed
Jul 28, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Thank you for getting this all up to date with the latest state of the repo, I greatly appreciate this. The adapter looks great, just some small nits on testing & other small configuration things that can sometimes slip by the wayside. Once these are wrapped up I will get this merged and ready for the next release!
Added comment to clarify the lack of tokenizer support for vision models.
Removed unused head_dim assignment from hf_config.
Updated error message for clarity and added handling for last_hidden_state in output.
Update test to verify that the forward method returns a tensor instead of a raw HF output object. Adjust assertions to match the expected behavior after changes in bridge.py.
Removed deprecated tests for prepare_loading() in TestViTPrepareLoading.
Contributor
Author
|
All issues resolved, all formats checked. All tests passed, ready to merge. |
Collaborator
|
Great work @david-wei-01001! |
Contributor
Author
|
Thank you very much! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Extended model_bridge of TransformerLens to support visual encoders (ViT, Deit)
Thoroughly tested:
unit test in
tests/unit/model_bridge/supported_architectures/test_vit_adapter.pyintergration test in
tests/integration/model_bridge/test_vit_adapter.pycan be tested as:
cd TransformerLens uv sync uv run pytest tests/unit/model_bridge/supported_architectures/test_vit_adapter.py -v uv run pytest tests/integration/model_bridge/test_vit_adapter.py -vFixes # (issue)
Type of change
Please delete options that are not relevant.
Screenshots
Please attach before and after screenshots of the change if applicable.
Checklist: