Install with pip install torch_vit
and load a pretrained ViT with:
from torch_vit import ViT
model = ViT('B_16_imagenet1k', pretrained=True)
This repository contains an PyTorch reimplementation of the Visual Transformer architecture from Google, along with pre-trained models and and wrapper modules and examples. Some other variations will developed in next releases
The goal of this implementation is to be simple, highly extensible, and easy to integrate into your own projects.
- At the moment, you can easily:
- Loading Pretrained ViT Models
- Evaluate on ImageClassification or Custom data
- Finetune ViT on Custom Dataset