Package: pretrained-vit-pytorch
TODOUrgentRewrite READMEImplement modelConvert pretrained weightsColab examplepipViT PyTorchQuickstartInstall withpip install vit_pytorchand load a pretrained ViT with:fromvit_pytorchimportViTmodel=ViT('B_16_imagenet1k',pretrained=True)Or find a Google Colab examplehere.OverviewThis repository contains an op-for-op PyTorch reimplementation of theVisual Transformerarchitecture fromGoogle, along with pre-trained models and examples.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:Load pretrained ViT modelsEvaluate on ImageNet or your own data(Upcoming features)Coming soon:Finetune ViT on your own datasetTrain ViT from scratch on ImageNet (1K)Export to ONNX for efficient inferenceTable of contentsAbout ViTAbout ViT-PyTorchInstallationUsageLoad pretrained modelsExample: ClassifyContributingAbout ViTVisual Transformers (ViT) are a straightforward application of thetransformer architectureto image classification. Even in computer vision, it seems, attention is all you need.The ViT architecture works as follows: (1) it considers an image as a 1-dimensional sequence of patches, (2) it prepends a classification token to the sequence, (3) it passes these patches through a transformer encoder (likeBERT), (4) it passes the first token of the output of the transformer through a small MLP to obtain the classification logits.
ViT is trained on a large-scale dataset (ImageNet-21k) with a huge amount of compute.About ViT-PyTorchViT-PyTorch is a PyTorch re-implementation of EfficientNet. It is consistent with theoriginal Jax implementation, so that it's easy to load Jax-pretrained weights.At the same time, we aim to make our PyTorch implementation as simple, flexible, and extensible as possible.InstallationInstall with pip:pipinstallvit_pytorchOr from source:gitclonehttps://github.com/lukemelas/ViT-PyTorchcdViT-Pytorch
pipinstall-e.UsageLoading pretrained modelsLoading a pretrained model is easy:fromvit_pytorchimportViTmodel=ViT('B_16_imagenet1k',pretrained=True)Details about the models are below:Name* Pretrained on *Finetuned on*Available? *B_16ImageNet-21k-✓B_16_imagenet1kImageNet-21kImageNet-1k✓B_32_imagenet1kImageNet-21kImageNet-1k✓L_16_imagenet1kImageNet-21kImageNet-1k✓L_32_imagenet1kImageNet-21kImageNet-1k✓Custom ViTLoading custom configurations is just as easy:fromvit_pytorchimportViT# The following is equivalent to ViT('B_16')config=dict(hidden_size=512,num_heads=8,num_layers=6)model=ViT.from_config(config)Example: ClassificationBelow is a simple, complete example. It may also be found as a Jupyter notebook inexamples/simpleor as aColab Notebook.importjsonfromPILimportImageimporttorchfromtorchvisionimporttransforms# Load ViTfromvit_pytorchimportViTmodel=ViT('B_16_imagenet1k',pretrained=True)model.eval()# Load image# NOTE: Assumes an image `img.jpg` exists in the current directoryimg=transforms.Compose([transforms.Resize((384,384)),transforms.ToTensor(),transforms.Normalize(0.5,0.5),])(Image.open('img.jpg')).unsqueeze(0)print(img.shape)# torch.Size([1, 3, 384, 384])# Classifywithtorch.no_grad():outputs=model(img)print(outputs.shape)# (1, 1000)ImageNetSeeexamples/imagenetfor details about evaluating on ImageNet.CreditOther great repositories with this model include:Ross Wightman's repoPhil Wang's repoContributingIf you find a bug, create a GitHub issue, or even better, submit a pull request. Similarly, if you have questions, simply post them as GitHub issues.I look forward to seeing what the community does with these models!