Refactor IJEPA to use timm.

[radiradev]

Changes

This PR adresses #1367. I have refactored IJEPA to use timm. I have tried to stay closer to the original implementation, and also added typing. There might be some structural changes needed - for instance the apply_masks function should probably be moved to utils? Any suggestions on how to improve this are welcome.

Also like the MAE timm implementation, I have created a separate file, instead of directly replacing the torchvision implementation. I assume once this is benchmarked the plan would be to completely replace it.

How was it tested?

Unit tests for the predictor, encoder and backbone classes.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 85.61%. Comparing base (78f59fc) to head (9708f36).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1612      +/-   ##
==========================================
+ Coverage   85.49%   85.61%   +0.11%     
==========================================
  Files         147      148       +1     
  Lines        6281     6333      +52     
==========================================
+ Hits         5370     5422      +52     
  Misses        911      911              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[guarin]

Hi! Thanks a lot for this extensive PR, it looks really well made!

Looking at the code I see many parallels to our MaskedVisionTransformer implementation. I didn't have time yet to go through the full PR but I have a suspicion that IJEPABackboneTIMM and IJEPAEncoderTIMM are compatible with our implementation of the MaskedVisionTransformer (see here and here). Do you think it would be possible to use MaskedVisionTransformer directly instead of IJEPABackboneTIMM or am I missing something? If possible this simplify the code a lot.

What I imagine is something like this:

target_encoder = MaskedVisionTransformer()
context_encoder = MaskedVisionTransformer()
predictor = IJEPAPredictorTIMM()

# This encodes all patches.
target = target_encoder.encode(images)
# This encodes only the unmasked context patches.
context = context_encoder.encode(images, idx_keep=idx_keep)

prediction = predictor(context, masks_x, masks)
target = get_targets_at_masks(target, masks_x, masks)
loss(predictions, target)

[radiradev]

Hi @guarin. I think you are correct and we can reuse the MaskedVisionTransformer making this PR a lot smaller. I have updated it with your suggested changes. Also added drop_path_rate for IJepaPredictor.

[guarin]

Awesome, thanks so much!

[guarin]

Created two follow-up issues:

• Update I-JEPA example to use timm #1614
• Add MaskedVisionTransformerDecoder #1615

Closes: #1367