vit_small_patch16_384.augreg_in21k_ft_in1k - 增强的视觉转换器模型及其在图像分类中的应用

vit_small_patch16_384.augreg_in21k_ft_in1k项目介绍

项目背景

vit_small_patch16_384.augreg_in21k_ft_in1k是一个用于图像分类的Vision Transformer (ViT)模型。该模型最初由论文作者在JAX上训练，并通过Ross Wightman移植到PyTorch框架中使用。训练分为两个阶段：首先在ImageNet-21k大数据集上进行预训练，然后在更小的ImageNet-1k数据集上进行微调。在微调的过程中，增加了数据增强和正则化技术，以提升模型的泛化能力。

模型详情

类型: 图像分类 / 特征提取骨干
参数数量: 22.2百万
GMACs: 12.4
激活数: 24.2百万
图像尺寸: 384 x 384

该模型在论文《How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers》和《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》中进行了详细描述。

数据集

预训练数据集: ImageNet-21k
微调数据集: ImageNet-1k

模型使用

图像分类

用户可以通过以下代码在图像上应用图像分类任务：

from urllib.request import urlopen
from PIL import Image
import timm

img = Image.open(urlopen(
    'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
))

model = timm.create_model('vit_small_patch16_384.augreg_in21k_ft_in1k', pretrained=True)
model = model.eval()

# 获取模型特定的变换（标准化、调整尺寸）
data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)

output = model(transforms(img).unsqueeze(0))  # 将单张图像扩展为批量大小为1

top5_probabilities, top5_class_indices = torch.topk(output.softmax(dim=1) * 100, k=5)

图像嵌入

模型也可以用于提取图像嵌入，适用于特征分析等任务：

from urllib.request import urlopen
from PIL import Image
import timm

img = Image.open(urlopen(
    'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
))

model = timm.create_model(
    'vit_small_patch16_384.augreg_in21k_ft_in1k',
    pretrained=True,
    num_classes=0,  # 去除分类器nn.Linear
)
model = model.eval()

# 获取模型特定的变换
data_config = timm.data.resolve_model_data_config(model)
transforms = timm.data.create_transform(**data_config, is_training=False)

output = model(transforms(img).unsqueeze(0))  # 输出为形状为(batch_size, num_features)的张量

模型对比

用户可通过timm模型结果页面探索更多关于该模型的数据集和运行时间指标的信息。

参考文献

本文涉及的模型和算法的详细信息可参阅以下参考文献：

Steiner, Andreas等人. "How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers." arXiv preprint arXiv:2106.10270, 2021.
Dosovitskiy, Alexey等人. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR, 2021.
Wightman, Ross. "PyTorch Image Models." GitHub, 2019.