ViT_pytorch 图片分类

以下是两种使用ViT_pytorch进行图片分类的例子：

1. 使用SimpleViT模型进行图片分类：

import torch
from vit_pytorch import SimpleViT

# 创建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)

2. 使用MobileViT模型进行图片分类：

import torch
from vit_pytorch.mobile_vit import MobileViT

# 创建MobileViT模型
mbvit_xs = MobileViT(
    image_size = (256, 256),
    dims = [96, 120, 144],
    channels = [16, 32, 48, 48, 64, 64, 80, 80, 96, 96, 384],
    num_classes = 1000
)

# 输入图片数据
img = torch.randn(1, 3, 256, 256)

# 进行预测
pred = mbvit_xs(img) # 输出：(1, 1000)