[FEAT] Add MobileViT v1 & v2

[yassineAlouini]

🚀 The feature

As described in the RFC "Batteries includes, phase 3", I am working on adding MobileViT v1 and v2 inspired by the following code repos/snippets:

• https://github.com/chinhsuanwu/mobilevit-pytorch
• https://github.com/apple/ml-cvnets/blob/main/cvnets/models/classification/mobilevit_v2.py

The original paper can be found here.

Motivation, pitch

This has been decided in the RFC.

Alternatives

No response

Additional context

No response

cc @datumbox

[datumbox]

Looks great @yassineAlouini. It would be great if we get this implementation.

Please have a read at #5319 where we document some best practices for model authoring. Also to avoid licensing problems, let's do a from scratch implementation.

[yassineAlouini]

Perfect, I will work on this today but mostly next week and the week after. Will let you know how my progress goes. 👌

[yassineAlouini]

I have started the implementation. Seems like a big chunk but excited to do it. 👌

I have found this huggingface implementation, could be useful as another inspiration: https://huggingface.co/docs/transformers/main/model_doc/mobilevit.

[EDIT] It looks like this is a wrapper around the https://github.com/apple/ml-cvnets implementation. 👌

[datumbox]

Hi @yassineAlouini. Just wanted to touch base on the implementation. Any blockers or need help?

[yassineAlouini]

Hello @datumbox, thanks for checking. So far, so good. It is taking a bit longer since I only had one day of working on it and it is paused for now but might work during the weekends and nights.

Do you expect a date for finishing? 🤔

[datumbox]

hey @yassineAlouini, sounds good. Thanks for the work. There are absolutely no deadlines on our side; just checking that everything goes smoothly and that you don't have a blocker. Let me know if you need anything :)

[yassineAlouini]

Some update @datumbox: I will have some free time for the upcoming few days and should make some progress. Will let you know how it goes. 👌

[yassineAlouini]

By the way, what is the PyTorch and TorchVision policies for the usage of einops? 🤔

[datumbox]

@yassineAlouini So far we don't have a model using this. Is there a specific use-case in MobileViT that can't be done otherwise?

[yassineAlouini]

I don't think it is irreplaceable, just wanted to check what is the best practice in torchvision. 👌
I will code everything using PyTorch and existing TorchVision code.

[yassineAlouini]

One additional question regarding the TransformerEncoder: should I reimplement it or should I re-use the one from vision_transformer.py (i.e. EncoderBlock)? I was planning to copy-paste the code first, adapt it and then maybe later refactor. What do you think @datumbox?

[datumbox]

@yassineAlouini Makes sense. Let's start by copy-pasting and modifying and see what changes are needed. Then we can decide whether sharing components is worth it. :)