Initial Vision Transformer architecture with MAE decoder

[weiji14]

What I am changing

• Initializing the neural network architecture, with a Vision Transformer (ViT) B/32 backbone and Masked Autoencoder (MAE) decoder

How I did it

• ViT backbone/decoder architecture is from HuggingFace transformers
  • Loosely using the ViT B/32 model, but with 12 channels instead of 3.

Note:

• Have consider torchvision/lightly (see fb8d06b), but changing input channels from 3 to 12 required a lot of hacking
  • Torchvision's ViT B/32 model - https://pytorch.org/vision/0.15/models/vision_transformer.html
  • MAE decoder and utils from Lightly (a self-supervised learning library)
    • Adapted from https://docs.lightly.ai/self-supervised-learning/examples/mae.html

TODO:

• Install transformers dependency
• Initialize model architecture backbone/decoder layers
• Setup training_step and forward pass
• Add unit tests
• Document model architecture in src/README.md

How you can test it

• Run python trainer.py fit --trainer.max_epochs=10 locally

Related Issues

Working towards #3

References:

• He, K., Chen, X., Xie, S., Li, Y., Dollar, P., & Girshick, R. (2022). Masked Autoencoders Are Scalable Vision Learners. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 15979–15988. https://doi.org/10.1109/CVPR52688.2022.01553

[srmsoumya]

The implementation looks good to me & we have enough options to modify in the MAE & ViT backbone.
Let us use this model for current sprint, next week we need to add options to:

1. Add embeddings for time, lat/lon, channels & position
2. Implement different masking strategy like random masking, grouped channel/time masking.
3. Add support for different backbones like SWIN or FlexiVIT

[weiji14]

Thanks @srm, I see you're starting the spatiotemporal embedding work on GeoViT at #47, and we can work on the masking strategy and different model backbones later too. I've renamed model.py to model_vit.py in case we want to have other model_*.py files. Will merge this into the main branch now.