jax-vision-transformers

Vision transformers with JAX & Flax

Installation

pip install jvt

Usage example

from jvt import DeiT

• ViT https://arxiv.org/abs/2010.11929
• MAE https://arxiv.org/abs/2111.06377
• DeiT https://arxiv.org/abs/2012.12877
• LeViT https://arxiv.org/abs/2104.01136
• ConvPass https://arxiv.org/abs/2207.07039