Skip to main content

Vision Transformer (ViT) - Pytorch

Project description

Pytorch Implementation of ViT

Original Paper link: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale(Alexey Dosovitskiy et al.)

Install

$ pip install vit-pytorch-implementation

#Usage:

import torch
from vit_pytorch import lilViT

v = lilViT(
                 img_size=224, 
                 in_channels=3,
                 patch_size=16, 
                 num_transformer_layers=12,
                 embedding_dim=768,
                 mlp_size=3072,
                 num_heads=12, 
                 attn_dropout=0,
                 mlp_dropout=0.1,
                 embedding_dropout=0.1,
                 num_classes=1000
)

img = torch.randn(1, 3, 224, 224)

preds = v(img) # (1, 1000)
preds.shape

Parameters

  • img_size: int.
    Image resolution. Default=224(224x224)
  • in_channels: int.
    Image channels. Default 3
  • patch_size: int.
    Size of patches. image_size must be divisible by patch_size.
    The number of patches is: n = (image_size // patch_size) ** 2 and n must be greater than 16. Default 16
  • num_transformer_layers: int.
    Depth(number of transformer blocks). Default 12
  • embedding_dim: int.
    Embedding dimension. Default 768
  • mlp_size: int.
    MLP size. Default 3072
  • num_heads: int.
    Number of heads in Multi-head Attention layer. Default 12
  • attn_dropout: float.
    Dropout for attention projection. Default 0
  • mlp_dropout: float
    Dropout for dense/MLP layers. Default 0.1
  • embedding_dropout: float.
    Dropout for patch and position embeddings.Default 0.1
  • num_classes: int.
    Number of classes to classify. Default 1000

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

vit-pytorch-implementation-1.0.2.tar.gz (3.8 kB view hashes)

Uploaded Source

Built Distribution

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page