Vision-LSTM (ViL)

[Project Page] [Paper] [Models] [Codebase Demo Video] [BibTeX]

Pytorch implementation and pre-trained models of Vision-LSTM (ViL), an adaption of xLSTM to computer vision.

License

This project is licensed under the MIT License, except the following folders/files, which are licensed under the AGPL-3.0 license:

src/vislstm/modules/xlstm
vision_lstm.py

Get started

This code-base supports simple usage of Vision-LSTM with an "architecture-only" implementation and also a full training pipeline.

Architecture only

The package vision_lstm provides a standalone implementation in the style of timm.

If you only need the model architecture, you can load it in a single line via torchhub or copy the vision_lstm into your own code-base.

# load ViL-T
model = torch.hub.load("nx-ai/vision-lstm", "VisionLSTM")
# load your own model
model = torch.hub.load(
    "nx-ai/vision-lstm", 
    "VisionLSTM",
    dim=192,  # latent dimension (192 for ViL-T)
    depth=24,  # how many ViL blocks
    patch_size=16,  # patch_size (results in 196 patches for 224x224 images)
    input_shape=(3, 224, 224),  # RGB images with resolution 224x224
    output_shape=(1000,),  # classifier with 1000 classes
    drop_path_rate=0.05,  # stochastic depth parameter
    stride=None,  # set to 8 for long-sequence fine-tuning
)

Full training pipeline (architecture, datasets, hyperparameters, ...)

If you want to train models with our code-base, follow the setup instructions from SETUP.md. To start runs, follow the instructions from RUN.md.

Pre-trained models

Pre-trained models on ImageNet-1K can be loaded via torchhub or directly downloaded from here.

# pre-trained models (Table 1, left)
model = torch.hub.load("nx-ai/vision-lstm", "vil-tiny")               # 77.3%
model = torch.hub.load("nx-ai/vision-lstm", "vil-tinyplus")           # 78.1%
model = torch.hub.load("nx-ai/vision-lstm", "vil-small")              # 80.7%
model = torch.hub.load("nx-ai/vision-lstm", "vil-smallplus")          # 80.9%
model = torch.hub.load("nx-ai/vision-lstm", "vil-base")               # 81.6%

# long-sequence fine-tuned models (Table 1, right)
model = torch.hub.load("nx-ai/vision-lstm", "vil-tinyplus-stride8")   # 80.0%
model = torch.hub.load("nx-ai/vision-lstm", "vil-smallplus-stride8")  # 82.2%
model = torch.hub.load("nx-ai/vision-lstm", "vil-base-stride8")       # 82.7%

# tiny models trained for only 400 epochs (Appendix A.2)
model = torch.hub.load("nx-ai/vision-lstm", "vil-tiny-e400")          # 76.1%
model = torch.hub.load("nx-ai/vision-lstm", "vil-tinyplus-e400")      # 77.2%

An example of how to use these models can be found in eval.py which evaluates the models on the ImageNet-1K validation set.

DeiT-III-T reimplementation models

Checkpoints for our reimplementation of DeiT-III-T are provided as raw checkpoint here and can be loaded from torchhub (the vision transformer implementation is based on KappaModules so you need to install it before loading a ViT checkpoint via torchhub by running pip install kappamodules==0.1.70).

model = torch.hub.load("nx-ai/vision-lstm", "deit3-tiny-e400")  # 75.6%
model = torch.hub.load("nx-ai/vision-lstm", "deit3-tiny")       # 76.2%

Other

This code-base is an improved version of the one used for MIM-Refiner for which there exists a demo video to explain various things.

Citation

If you like our work, please consider giving it a star ⭐ and cite us

@article{alkin2024visionlstm,
  title={Vision-LSTM: xLSTM as Generic Vision Backbone},
  author={Benedikt Alkin and Maximilian Beck and Korbinian P{\"o}ppel and Sepp Hochreiter and Johannes Brandstetter}
  journal={arXiv preprint arXiv:2406.04303},
  year={2024}
}

NX-AI / vision-lstm