Finish I-JEPA

[guarin]

Experimental support for I-JEPA was added in #1273

We should do some refactoring and testing before fully releasing the model.

Todo

• Refactor I-JEPA to use timm ViT #1367
• Verify that weights are correctly initialized and follow this
• Check if we can refactor IJEPAMaskCollator into a transform, comment.
• Move positional embedding functions (_get_1d_sincos_pos_embed_from_grid etc.) to lightly/models/utils and add reference to source
• Add missing docstrings or bring into correct format, check diff from I-JEPA #1273
• Add unit tests
• Finish pytorch lightning and distributed examples
• Add docs
• Add imagenet benchmark to lightly/benchmarks/imagenet/vit/ijepa.py
• Fix broken IJEPA Example #1712

[Natyren]

Hello @guarin . I can work on it. Or on part of issues.

[guarin]

That would great! Feel free to pick anything up that you are interested in. Let me know if you need help or some clarification.

[Natyren]

@guarin
I will take issues

• Try if we can get it working with timm ViT backbone. I would like to do this because torchvision ViT doesn't support stochastic path dropout, see comment.
• Verify that weights are correctly initialized and follow this
• Check if we can refactor IJEPAMaskCollator into a transform, comment.
• Move positional embedding functions (_get_1d_sincos_pos_embed_from_grid etc.) to lightly/models/utils and add reference to source
• Finish pytorch lightning and distributed examples
• Add docs
• Add imagenet benchmark to lightly/benchmarks/imagenet/vit/ijepa.py

Will work on them in next prs

[Natyren]

#1322 reference point to this issue

• Move positional embedding functions (_get_1d_sincos_pos_embed_from_grid etc.) to lightly/models/utils and add reference to source

[Natyren]

Next, i will work on implementation of this

• Check if we can refactor IJEPAMaskCollator into a transform, comment.

UPDATE: on pause