vit_small_r26_s32_384.augreg_in21k_ft_in1k - ResNet与Vision Transformer结合的图像分类模型解析

vit_small_r26_s32_384.augreg_in21k_ft_in1k项目介绍

项目背景

vit_small_r26_s32_384.augreg_in21k_ft_in1k是一个图像分类模型，它将ResNet与视觉Transformer（Vision Transformer, ViT）相结合。这一模型的设计和训练依赖于强大的数据集和精细的调整，以提高图像识别的精确度。模型最初由论文作者在JAX平台上进行训练，然后Ross Wightman将其移植到PyTorch中。

模型详情

模型类型：用于图像分类，亦可作为特征提取的骨干网络。
参数统计：
- 参数数量（百万）：36.5
- GMACs：10.2（计算复杂度衡量单位）
- 激活（百万）：27.7
- 输入图像大小：384 x 384像素
相关文献：
- 文章《How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers》探讨了如何通过数据扩增和正则化来优化ViT训练。
- 文章《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》介绍了Transformer在大规模图像识别中的应用。
使用数据集：
- 训练数据集：ImageNet-21k
- 微调数据集：ImageNet-1k

模型使用

图像分类

用户可以通过加载预训练模型来进行图像分类。该模型支持在标准图像输入上进行预测，并输出图像属于不同类别的概率，其中最高的5个类别概率通常被作为最终输出。

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

img = Image.open(urlopen('图像URL'))
model = timm.create_model('vit_small_r26_s32_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))
top5_probabilities, top5_class_indices = torch.topk(output.softmax(dim=1) * 100, k=5)

图像嵌入

除了分类之外，模型还可以用于提取图像嵌入（image embeddings），以供进一步的图像分析或下游任务使用。这种方法允许用户在不需要分类头的情况下获取潜在特征。

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

img = Image.open(urlopen('图像URL'))

model = timm.create_model(
    'vit_small_r26_s32_384.augreg_in21k_ft_in1k',
    pretrained=True,
    num_classes=0,  # 移除分类层
)
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))

模型比较

用户可以通过timm库的模型结果页面，浏览该模型在不同数据集和运行环境下的表现，以便进行合理的模型选择和应用。

推荐引用

在学术研究或项目中使用本模型时，建议引用以下文献，这些文献为模型的开发和性能提升提供了理论基础。

@article{steiner2021augreg,
  title={How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers},
  author={Steiner, Andreas et al.},
  journal={arXiv preprint arXiv:2106.10270},
  year={2021}
}

@article{dosovitskiy2020vit,
  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},
  author={Dosovitskiy, Alexey et al.},
  journal={ICLR},
  year={2021}
}

@misc{rw2019timm,
  author = {Ross Wightman},
  title = {PyTorch Image Models},
  year = {2019},
  doi = {10.5281/zenodo.4414861},
  url = {https://github.com/huggingface/pytorch-image-models}
}

这个项目的目标是将Transformer应用于图像领域，并探索其在图像分类任务中的有效性。通过整合丰富的数据集和先进的理论研究，该模型展示了在视觉任务上的显著性能。