kaicheng-yang0828 / vision-transformer-vit

Pytorch ViT for Image classification on the CIFAR10 dataset

Vit-ImageClassification

Introduction

This project uses ViT to perform image classification tasks on DATA set CIFAR10. The implement of Vit and pretrained weight are from https://github.com/asyml/vision-transformer-pytorch.

Installation

pytorch 1.7.1 python 3.7.3

Datasets

Download the CIFAR10 from http://www.cs.toronto.edu/~kriz/cifar.html or you can get it from https://pan.baidu.com/s/1ogAFopdVzswge2Aaru_lvw (code: k5v8), creat data floder and unzip the cifar-10-python.tar.gz under './data'.

Pre_trained model

You can download the pretrained file from https://pan.baidu.com/s/1CuUj-XIXwecxWMEcLoJzPg (code: ox9n), creat Vit_weights floder and pretrained file under ./Vit_weights

Train

python main.py 

Result

Base on the pretrained weight, after one epoch, we get 98.1 Accuracy

model	dataset	acc
ViT-B_16	CIFAR10	98.1

vision-transformer-vit's Issues

How long does it take to train?

Could you please describe you set up (GPU etc) as well? Thank you!

`total_steps` in train.py

There seems to be a problem in train.py.

total_steps = (len(trainloader) // args.batch_size + 1)  * args.epoches

len(train_loader) is already divided by batch_size. Change it to total_steps = len(trainloader) * args.epoches and it should work!