GitHub：ViT-pytorch相关学习-视觉分类方向-1_vit github

pip install vit-pytorch

import torch
from vit_pytorch import ViT
 
v = ViT(
    image_size = 256,     
    patch_size = 32,      
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)
 
img = torch.randn(1, 3, 256, 256)
 
preds = v(img) # (1, 1000)

import torch
from vit_pytorch import SimpleViT
 
v = SimpleViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048
)
 
img = torch.randn(1, 3, 256, 256)
 
preds = v(img) # (1, 1000)

import torch
from torchvision.models import resnet50
 
from vit_pytorch.distill import DistillableViT, DistillWrapper
 
teacher = resnet50(pretrained = True)
 
v = DistillableViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 8,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)
 
distiller = DistillWrapper(
    student = v,
    teacher = teacher,
    temperature = 3,           # temperature of distillation
    alpha = 0.5,               # trade between main loss and distillation loss
    hard = False               # whether to use soft or hard distillation
)
 
img = torch.randn(2, 3, 256, 256)
labels = torch.randint(0, 1000, (2,))
 
loss = distiller(img, labels)
loss.backward()
 
# after lots of training above ...
 
pred = v(img) # (2, 1000)

v = v.to_vit()
type(v) # <class 'vit_pytorch.vit_pytorch.ViT'>

import torch
from vit_pytorch.deepvit import DeepViT
 
v = DeepViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)
 
img = torch.randn(1, 3, 256, 256)
 
preds = v(img) # (1, 1000)

import torch
from vit_pytorch.cait import CaiT
 
v = CaiT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 12,             # depth of transformer for patch to patch attention only
    cls_depth = 2,          # depth of cross attention of CLS tokens to patch
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1,
    layer_dropout = 0.05    # randomly dropout 5% of the layers
)
 
img = torch.randn(1, 3, 256, 256)
 
preds = v(img) # (1, 1000)

import torch
from vit_pytorch.t2t import T2TViT
 
v = T2TViT(
    dim = 512,
    image_size = 224,
    depth = 5,
    heads = 8,
    mlp_dim = 512,
    num_classes = 1000,
    t2t_layers = ((7, 4), (3, 2), (3, 2)) # tuples of the kernel size and stride of each consecutive layers of the initial token to token module
)
 
img = torch.randn(1, 3, 224, 224)
 
preds = v(img) # (1, 1000)

import torch
from vit_pytorch.cct import CCT
 
cct = CCT(
    img_size = (224, 448),
    embedding_dim = 384,
    n_conv_layers = 2,
    kernel_size = 7,
    stride = 2,
    padding = 3,
    pooling_kernel_size = 3,
    pooling_stride = 2,
    pooling_padding = 1,
    num_layers = 14,
    num_heads = 6,
    mlp_radio = 3.,
    num_classes = 1000,
    positional_embedding = 'learnable', # ['sine', 'learnable', 'none']
)
 
img = torch.randn(1, 3, 224, 448)
pred = cct(img) # (1, 1000)

import torch
from vit_pytorch.cct import cct_14
 
cct = cct_14(
    img_size = 224,
    n_conv_layers = 1,
    kernel_size = 7,
    stride = 2,
    padding = 3,
    pooling_kernel_size = 3,
    pooling_stride = 2,
    pooling_padding = 1,
    num_classes = 1000,
    positional_embedding = 'learnable', # ['sine', 'learnable', 'none']
)

import torch
from vit_pytorch.cross_vit import CrossViT
 
v = CrossViT(
    image_size = 256,
    num_classes = 1000,
    depth = 4,               # number of multi-scale encoding blocks
    sm_dim = 192,            # high res dimension
    sm_patch_size = 16,      # high res patch size (should be smaller than lg_patch_size)
    sm_enc_depth = 2,        # high res depth
    sm_enc_heads = 8,        # high res heads
    sm_enc_mlp_dim = 2048,   # high res feedforward dimension
    lg_dim = 384,            # low res dimension
    lg_patch_size = 64,      # low res patch size
    lg_enc_depth = 3,        # low res depth
    lg_enc_heads = 8,        # low res heads
    lg_enc_mlp_dim = 2048,   # low res feedforward dimensions
    cross_attn_depth = 2,    # cross attention rounds
    cross_attn_heads = 8,    # cross attention heads
    dropout = 0.1,
    emb_dropout = 0.1
)
 
img = torch.randn(1, 3, 256, 256)
 
pred = v(img) # (1, 1000)

import torch
from vit_pytorch.pit import PiT
 
v = PiT(
    image_size = 224,
    patch_size = 14,
    dim = 256,
    num_classes = 1000,
    depth = (3, 3, 3),     # list of depths, indicating the number of rounds of each stage before a downsample
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)
 
# forward pass now returns predictions and the attention maps
 
img = torch.randn(1, 3, 224, 224)
 
preds = v(img) # (1, 1000)