标签:

• 图像特征提取
• 鸟类识别
• pytorch 库名称: birder 许可证: apache-2.0

vit_reg4_b16_mim 模型卡

一个采用掩码图像建模(MIM)预训练的ViT reg4图像编码器。该模型未针对特定分类任务进行微调，旨在作为通用特征提取器或用于下游任务（如目标检测、分割或自定义分类）的骨干网络。

模型详情

• 模型类型: 图像编码器

• 模型统计:

  • 参数量(M): 85.8
  • 输入图像尺寸: 224 x 224

• 数据集: 在约1100万张多样化图像上训练，包括:

  • iNaturalist 2021 (~330万)
  • WebVision-2.0 (~150万随机子集)
  • imagenet-w21-webp-wds (~100万随机子集)
  • SA-1B (~22万随机子集，选自20个分块)
  • COCO (~12万)
  • NABirds (~4.8万)
  • Birdsnap v1.1 (~4.4万)
  • CUB-200 2011 (~1.8万)
  • Birder专有数据集 (~500万)

• 相关论文:

  • 《一张图像相当于16x16个单词：大规模图像识别的Transformer模型》: https://arxiv.org/abs/2010.11929
  • 《视觉Transformer需要寄存器》: https://arxiv.org/abs/2309.16588
  • 《掩码自编码器是可扩展的视觉学习器》: https://arxiv.org/abs/2111.06377

模型使用

图像嵌入提取

import torch
import birder
from PIL import Image

(net, model_info) = birder.load_pretrained_model("vit_reg4_b16_mim_300", inference=True)

# 获取模型训练时使用的图像尺寸
size = birder.get_size_from_signature(model_info.signature)

# 创建推理转换流程
transform = birder.classification_transform(size, model_info.rgb_stats)

image = Image.open("图片路径.jpeg")
input_tensor = transform(image).unsqueeze(dim=0)
with torch.inference_mode():
    embedding = net.embedding(input_tensor)
    # embedding是形状为(1, 768)的张量

或使用load_model_with_cfg函数

import torch
import birder
from PIL import Image

# 需先下载模型文件
(net, cfg) = birder.load_model_with_cfg("models/vit_reg4_b16_mim.json", "models/vit_reg4_b16_mim_300.pt")
net.eval()

# 获取模型训练时使用的图像尺寸
size = birder.get_size_from_signature(cfg["signature"])

# 创建推理转换流程
transform = birder.classification_transform(size, cfg["rgb_stats"])

image = Image.open("图片路径.jpeg")
input_tensor = transform(image).unsqueeze(dim=0)
with torch.inference_mode():
    embedding = net.embedding(input_tensor)
    # embedding是形状为(1, embedding_size)的张量

引用文献

@misc{dosovitskiy2021imageworth16x16words,
      title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale}, 
      author={Alexey Dosovitskiy and Lucas Beyer and Alexander Kolesnikov and Dirk Weissenborn and Xiaohua Zhai and Thomas Unterthiner and Mostafa Dehghani and Matthias Minderer and Georg Heigold and Sylvain Gelly and Jakob Uszkoreit and Neil Houlsby},
      year={2021},
      eprint={2010.11929},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2010.11929}, 
}

@misc{darcet2024visiontransformersneedregisters,
      title={Vision Transformers Need Registers}, 
      author={Timothée Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski},
      year={2024},
      eprint={2309.16588},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2309.16588}, 
}

@misc{he2021maskedautoencodersscalablevision,
      title={Masked Autoencoders Are Scalable Vision Learners}, 
      author={Kaiming He and Xinlei Chen and Saining Xie and Yanghao Li and Piotr Dollár and Ross Girshick},
      year={2021},
      eprint={2111.06377},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2111.06377}, 
}