TFAugViT model is the tensorflow implementation of the
AugViT: Augmented Shortcuts for Vision Transformers by Yehui Tang, Kai Han, Chang Xu, An Xiao, Yiping Deng, Chao Xu and Yunhe Wang,
and first released in
this repository.
Aug-ViT inserts additional paths with learnable parameters in parallel on the original shortcuts for alleviating the feature collapse. The block-circulant projection is used to implement augmented shortcut, which brings negligible increase of computational cost.
This model can be used for image classification tasks and easily be fine-tuned to suite your purpose of use.
Here is how to use this model to classify an image into one of the 1,000 ImageNet classes:
1from transformers import TFAutoModelForImageClassification
2from PIL import Image
3import requests
4
5url = "http://images.cocodataset.org/val2017/000000039769.jpg"
6image = Image.open(requests.get(url, stream=True).raw)
7
8model = TFAutoModelForImageClassification.from_pretrained("tensorgirl/TFaugvit",trust_remote_code=True)
9
10outputs = model({'pixel_values':image})
11
12
13# model predicts one of the 1000 ImageNet classes
14predicted_class_idx = outputs.argmax(-1)
The TFAugViT model is trained on
ImageNet-1k, a dataset consisting of 1 million images and 1,000 classes.
Due to the use of einops library you cannot use the model,fit() directly on this model, you will have to either write a custom training loop by passing the inputs as shown above or you can wrap the model in a functional model of keras and specify the batch_size beforehand.
If you want to train the model on some other data then either resize the images to 224x224 or change the model config image_size to suit your requirements.
1@inproceedings{aug-vit tf,
2title = {AugViT: Augmented Shortcuts for Vision Transformers},
3author = {Yehui Tang, Kai Han, Chang Xu, An Xiao, Yiping Deng, Chao Xu and Yunhe Wang},
4year = {2021},
5URL = {https://arxiv.org/abs/2106.15941}
6}
