Add I-JEPA #33125

jmtzt · 2024-08-26T14:40:37Z

What does this PR do?

This PR adds I-JEPA.

To-Do's:

convert remaining checkpoints ijepa_vith14_22k, ijepa_vith16_1k, ijepa_vitg16_22k.
transfer checkpoints to the meta org

ArthurZucker · 2024-08-27T13:37:38Z

cc @amyeroberts and @qubvel

amyeroberts

Thanks for adding this model @jmtzt!

Overall it's looking great. Did an initial review outlining some small things to update.

Before merging, we'll need to make sure the slow model tests are passing. To trigger these, you'll need to push a commit (empty or otherwise) with the message [run_slow] ijpea. Me or another person at HF will then need to approve the workflow

docs/source/en/model_doc/ijepa.md

src/transformers/models/ijepa/__init__.py

src/transformers/models/ijepa/configuration_ijepa.py

tests/models/ijepa/test_modeling_ijepa.py

src/transformers/models/ijepa/modeling_ijepa.py

jmtzt · 2024-08-28T08:51:24Z

Thanks for reviewing the code, @amyeroberts! :)

I've just pushed the adjustments, and now the fx support seems to be working fine.

amyeroberts

Thanks for iterating - looks great!

Just a small nit. It seems the slow model tests didn't run -- not really sure what happened there. Could you try with [run-slow] ijepa instead of [run_slow] ijepa?

After than, just a merge conflict to resolve and we're good to go!

src/transformers/models/ijepa/modeling_ijepa.py

jmtzt · 2024-08-30T18:30:26Z

Thanks again for checking @amyeroberts :) I've just pushed the adjustments, however some tests appear to be failing due to some OS errors inside CircleCI, do you know why's that?

amyeroberts · 2024-09-02T11:03:14Z

Thanks for pushing again! Hmmmm - I'm not sure why the slow model tests aren't being picked up or run here cc @ydshieh

amyeroberts · 2024-09-02T15:30:45Z

Yih-Dar found the issue - the actions won't be triggered whilst there are merge conflicts in the PR. Could you resolve these, then push another [run-slow] ijepa commit? This should hopefully run the workflow!

jmtzt · 2024-11-18T19:08:03Z

the original I-JEPA model doesn't have the pooling layer, so I think to get around this we might need to default the add_pooling_layer to False in its initialization, and modify this snippet accordingly to get the last hidden states rather than the pooler_output

qubvel · 2024-11-18T19:34:18Z

Ok, sounds good 👍

qubvel

Hi @jmtzt! Seems like almost everything is fine! Thanks for correcting snippets in docs and model cards! Snippets work fine on my side, I also did an experiment fine-tuning classification model - converges really good.

A few final nits regarding docstrings/constants and a suggestion regarding classification head and I will pass it to a core maintainers review, thank you for the great job!

src/transformers/models/ijepa/modeling_ijepa.py

tests/models/ijepa/test_modeling_ijepa.py

docs/source/en/model_doc/ijepa.md

jmtzt · 2024-11-19T14:26:34Z

hi @qubvel, thanks for the support and reviewing the PR! Just pushed the suggested changes, let me know if it's alright now.

qubvel

Thanks for addressing all the comments quickly! IMO it's ready for the next review

P.S. It may take a while since there are a lot of PRs for the final review in a line 🤗

qubvel · 2024-11-19T14:40:01Z

src/transformers/models/ijepa/modeling_ijepa.py

+        checkpoint=_IMAGE_CLASS_CHECKPOINT,
+        output_type=ImageClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+        expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT,


I suppose we can delete _IMAGE_CLASS_CHECKPOINT and _IMAGE_CLASS_EXPECTED_OUTPUT because we don't really have pretrained checkpoint for such a model

qubvel · 2024-11-19T14:42:02Z

src/transformers/models/ijepa/modeling_ijepa.py

+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+_EXPECTED_OUTPUT_SHAPE = [1, 197, 768]


hmm, it is not overwritten in modeling, while it exists in a modular file.. maybe a modular converter issue

in the modular_model_converter.py implementation, line 472, we have something like:

These top-level variables will always use the value in the modular_xxx.py file
ASSIGNMENTS_TO_KEEP = {"_CHECKPOINT_FOR_DOC", }

ok, maybe we have to add "_EXPECTED_OUTPUT_SHAPE" too

qubvel · 2024-11-19T14:53:32Z

@ArthurZucker please review whenever you have bandwidth! The model is similar to ViT, so the Modular is used here to remove the CLS token.

Checkpoints can be found here:
https://huggingface.co/jmtzt

(we will need to transfer them to facebook org as soon as we ensure the code is in the final stage + rename all occurrences in code and model cards)

ArthurZucker

Thanks a lot for the PR! Modular makes it easy to understand: basically VIT but no cls token embedding right? (checked that no models in transformers already has this!)

src/transformers/models/auto/configuration_auto.py

src/transformers/models/ijepa/__init__.py

src/transformers/models/ijepa/modular_ijepa.py

ArthurZucker · 2024-11-20T17:59:43Z

src/transformers/models/ijepa/modular_ijepa.py

+        num_patches = self.patch_embeddings.num_patches
+        self.position_embeddings = nn.Parameter(torch.randn(1, num_patches, config.hidden_size))
+
+    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:


any reason why this was resolved?

ArthurZucker · 2024-11-20T18:04:53Z

src/transformers/models/ijepa/modular_ijepa.py

+        self.embeddings = IJepaEmbeddings(config, use_mask_token=use_mask_token)
+        self.encoder = IJepaEncoder(config)
+        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.pooler = IJepaPooler(config) if add_pooling_layer else None
+        # Initialize weights and apply final processing
+        self.post_init()


Suggested change

self.embeddings = IJepaEmbeddings(config, use_mask_token=use_mask_token)

self.encoder = IJepaEncoder(config)

self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

self.pooler = IJepaPooler(config) if add_pooling_layer else None

# Initialize weights and apply final processing

self.post_init()

self.embeddings = IJepaEmbeddings(config, use_mask_token=use_mask_token)

this should be the only thing you have to change as all the other classes are the exact same

src/transformers/models/ijepa/modular_ijepa.py

ArthurZucker · 2024-11-20T18:06:00Z

tests/models/ijepa/test_modeling_ijepa.py

IMO you should use either copied form or inheritance for the tests as the model is so similar to VIT models. The most important test to add is an integration tests!

Example of inheritance in tests with Gemma2 !

There are already integration tests implemented, quite similar to the ViT ones actually, the only one missing is the test_inference_image_classification_head, and IMO it doesn't make sense to have this, as we don't have a checkpoint for image classification tasks.

also for your comment on the interpolate_pos_encoding, my implementation is different, to account for the lack of a CLS token...

ArthurZucker · 2024-11-20T18:24:29Z

src/transformers/models/ijepa/modular_ijepa.py

+        num_patches = self.patch_embeddings.num_patches
+        self.position_embeddings = nn.Parameter(torch.randn(1, num_patches, config.hidden_size))
+
+    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:


We should not even have to right it in the modular as it should be inherited!

ArthurZucker · 2024-11-25T14:57:53Z

Feel free to ping me again for another review! 🫡

jmtzt · 2024-11-30T12:10:59Z

Feel free to ping me again for another review! 🫡

@ArthurZucker I've pushed the adjustments according to your comments, let me know if anything is missing. Thanks! :)

ArthurZucker

Very nice! Camel casing is a bit wrong but would look super ugly otherwise!

qubvel · 2024-12-05T16:12:40Z

@ArthurZucker can you transfer checkpoints to Facebook org? Also, code snippets need to be adjusted.
Checkpoints are here: https://huggingface.co/jmtzt

qubvel · 2024-12-05T18:36:00Z

@jmtzt congratulations on the model merged 🎉 was glad to collaborate with you on this!
Can you please confirm if we can transfer checkpoints to Facebook org from your account?

jmtzt · 2024-12-05T18:39:47Z

thanks for the support @qubvel and @ArthurZucker :)

Sure, go ahead!

* first draft * add IJepaEmbeddings class * fix copy-from for IJepa model * add weight conversion script * update attention class names in IJepa model * style changes * Add push_to_hub option to convert_ijepa_checkpoint function * add initial tests for I-JEPA * minor style changes to conversion script * make fixup related * rename conversion script * Add I-JEPA to sdpa docs * minor fixes * adjust conversion script * update conversion script * adjust sdpa docs * [run_slow] ijepa * [run-slow] ijepa * [run-slow] ijepa * [run-slow] ijepa * [run-slow] ijepa * [run-slow] ijepa * formatting issues * adjust modeling to modular code * add IJepaModel to objects to ignore in docstring checks * [run-slow] ijepa * fix formatting issues * add usage instruction snippet to docs * change pos encoding, add checkpoint for doc * add verify logits for all models * [run-slow] ijepa * update docs to include image feature extraction instructions * remove pooling layer from IJepaModel in image classification class * [run-slow] ijepa * remove pooling layer from IJepaModel constructor * update docs * [run-slow] ijepa * [run-slow] ijepa * small changes * [run-slow] ijepa * style adjustments * update copyright in init file * adjust modular ijepa * [run-slow] ijepa

jmtzt added 13 commits August 22, 2024 14:32

first draft

3f53abd

add IJepaEmbeddings class

b9d7c03

fix copy-from for IJepa model

7af8961

add weight conversion script

a4c8eec

update attention class names in IJepa model

bf70f98

style changes

64f2208

Add push_to_hub option to convert_ijepa_checkpoint function

1dd4e7d

add initial tests for I-JEPA

9826f99

minor style changes to conversion script

d78e468

make fixup related

7a64b83

rename conversion script

66773ee

Add I-JEPA to sdpa docs

9b7e8b4

Merge branch 'huggingface:main' into add_ijepa

edd2ac9

ArthurZucker added the New model label Aug 27, 2024

amyeroberts added the run-slow label Aug 27, 2024

amyeroberts reviewed Aug 27, 2024

View reviewed changes

jmtzt added 5 commits August 28, 2024 10:22

minor fixes

40cf528

adjust conversion script

2bae64a

update conversion script

4ccf28c

adjust sdpa docs

851ed7e

[run_slow] ijepa

b7a027c

jmtzt requested a review from amyeroberts August 30, 2024 17:08

amyeroberts approved these changes Aug 30, 2024

View reviewed changes

src/transformers/models/ijepa/modeling_ijepa.py Outdated Show resolved Hide resolved

[run-slow] ijepa

552e800

[run-slow] ijepa

f2f7eb8

jmtzt added 4 commits November 18, 2024 21:17

remove pooling layer from IJepaModel constructor

db79009

update docs

57e5407

[run-slow] ijepa

8236816

[run-slow] ijepa

ce6499f

qubvel reviewed Nov 19, 2024

View reviewed changes

jmtzt added 2 commits November 19, 2024 15:16

small changes

81a6e66

[run-slow] ijepa

7a0fc39

qubvel approved these changes Nov 19, 2024

View reviewed changes

qubvel requested a review from ArthurZucker November 19, 2024 14:47

ArthurZucker reviewed Nov 20, 2024

View reviewed changes

ArthurZucker removed the request for review from amyeroberts November 20, 2024 18:27

jmtzt added 4 commits November 26, 2024 16:33

style adjustments

37a38f9

update copyright in init file

491d5a5

adjust modular ijepa

2afaba0

[run-slow] ijepa

db4dfc0

jmtzt requested a review from ArthurZucker November 26, 2024 15:35

ArthurZucker approved these changes Dec 5, 2024

View reviewed changes

ArthurZucker merged commit 50189e3 into huggingface:main Dec 5, 2024
22 checks passed

xenova mentioned this pull request Dec 5, 2024

Add ONNX export support for IJepa huggingface/optimum#2118

Draft

3 tasks

NielsRogge mentioned this pull request Dec 8, 2024

[I-JEPA] Update docs #35148

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add I-JEPA #33125

Add I-JEPA #33125

jmtzt commented Aug 26, 2024 •

edited

Loading

ArthurZucker commented Aug 27, 2024

amyeroberts left a comment

jmtzt commented Aug 28, 2024

amyeroberts left a comment

jmtzt commented Aug 30, 2024 •

edited

Loading

amyeroberts commented Sep 2, 2024

amyeroberts commented Sep 2, 2024

jmtzt commented Nov 18, 2024

qubvel commented Nov 18, 2024

qubvel left a comment

jmtzt commented Nov 19, 2024

qubvel left a comment •

edited

Loading

qubvel Nov 19, 2024

qubvel Nov 19, 2024

jmtzt Nov 19, 2024

qubvel Nov 19, 2024

qubvel commented Nov 19, 2024

ArthurZucker left a comment

ArthurZucker Nov 20, 2024

ArthurZucker Nov 20, 2024

ArthurZucker Nov 20, 2024

ArthurZucker Nov 20, 2024

jmtzt Nov 26, 2024

jmtzt Nov 26, 2024

ArthurZucker Nov 20, 2024

ArthurZucker commented Nov 25, 2024

jmtzt commented Nov 30, 2024

ArthurZucker left a comment

qubvel commented Dec 5, 2024

qubvel commented Dec 5, 2024 •

edited

Loading

jmtzt commented Dec 5, 2024

Add I-JEPA #33125

Add I-JEPA #33125

Conversation

jmtzt commented Aug 26, 2024 • edited Loading

What does this PR do?

ArthurZucker commented Aug 27, 2024

amyeroberts left a comment

Choose a reason for hiding this comment

jmtzt commented Aug 28, 2024

amyeroberts left a comment

Choose a reason for hiding this comment

jmtzt commented Aug 30, 2024 • edited Loading

amyeroberts commented Sep 2, 2024

amyeroberts commented Sep 2, 2024

jmtzt commented Nov 18, 2024

qubvel commented Nov 18, 2024

qubvel left a comment

Choose a reason for hiding this comment

jmtzt commented Nov 19, 2024

qubvel left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel commented Nov 19, 2024

ArthurZucker left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ArthurZucker commented Nov 25, 2024

jmtzt commented Nov 30, 2024

ArthurZucker left a comment

Choose a reason for hiding this comment

qubvel commented Dec 5, 2024

qubvel commented Dec 5, 2024 • edited Loading

jmtzt commented Dec 5, 2024

jmtzt commented Aug 26, 2024 •

edited

Loading

jmtzt commented Aug 30, 2024 •

edited

Loading

qubvel left a comment •

edited

Loading

qubvel commented Dec 5, 2024 •

edited

Loading