🚀 ONNX 轉換的 ViT(基礎大小模型)
本項目是 ViT-base 的 ONNX 轉換版本,帶有分類頭,可用於圖像分類任務。
🚀 快速開始
本項目的 Vision Transformer(ViT)模型在 ImageNet-21k(1400 萬張圖像,21843 個類別)上以 224x224 分辨率進行預訓練,並在 ImageNet 2012(100 萬張圖像,1000 個類別)上以相同分辨率進行微調。該模型由 Dosovitskiy 等人在論文 An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale 中提出,並首次在 此倉庫 發佈。不過,權重是從 Ross Wightman 的 timm 倉庫 轉換而來,他已將權重從 JAX 轉換為 PyTorch,在此對他表示感謝。
需要注意的是,發佈 ViT 的團隊並未為此模型撰寫模型卡片,此模型卡片由 Hugging Face 團隊編寫。
✨ 主要特性
- 基於 Vision Transformer(ViT)架構,在大規模圖像數據集上預訓練,具有強大的圖像特徵提取能力。
- 經過 ONNX 轉換,可在支持 ONNX 的推理引擎上高效運行。
- 適用於圖像分類任務,可對圖像進行準確的類別預測。
📦 安裝指南
文檔中未提及具體安裝步驟,可參考相關依賴庫的官方文檔進行安裝,如 transformers
、optimum
等。
💻 使用示例
基礎用法
以下是如何使用此模型將 COCO 2017 數據集中的圖像分類為 1000 個 ImageNet 類別之一的示例:
from transformers import AutoFeatureExtractor
from optimum.onnxruntime import ORTModelForImageClassification
from optimum.pipelines import pipeline
feature_extractor = AutoFeatureExtractor.from_pretrained("optimum/vit-base-patch16-224")
model = ORTModelForImageClassification.from_pretrained("optimum/vit-base-patch16-224")
onnx_img_classif = pipeline(
"image-classification", model=model, feature_extractor=feature_extractor
)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
pred = onnx_img_classif(url)
print("Top-5 predicted classes:", pred)
📚 詳細文檔
模型描述
Vision Transformer(ViT)是一種基於 Transformer 編碼器的模型(類似 BERT),它在大規模圖像數據集 ImageNet-21k 上以 224x224 像素的分辨率進行有監督的預訓練。之後,該模型在 ImageNet(也稱為 ILSVRC2012)數據集上進行微調,該數據集包含 100 萬張圖像和 1000 個類別,分辨率同樣為 224x224。
圖像以固定大小的圖像塊序列(分辨率為 16x16)的形式輸入到模型中,並進行線性嵌入。同時,會在序列開頭添加一個 [CLS] 標記,用於分類任務。在將序列輸入到 Transformer 編碼器的各層之前,還會添加絕對位置嵌入。
通過預訓練,模型學習到圖像的內部表示,可用於提取對下游任務有用的特徵。例如,如果有一個帶標籤的圖像數據集,可以在預訓練的編碼器之上放置一個線性層,訓練一個標準分類器。通常會在線性層放置在 [CLS] 標記之上,因為該標記的最後隱藏狀態可以看作是整個圖像的表示。
預期用途和侷限性
可以使用原始模型進行圖像分類。可在 模型中心 查找針對感興趣任務的微調版本。
🔧 技術細節
訓練數據
ViT 模型在 ImageNet-21k 上進行預訓練,該數據集包含 1400 萬張圖像和 21k 個類別;並在 ImageNet 上進行微調,該數據集包含 100 萬張圖像和 1k 個類別。
訓練過程
預處理
訓練/驗證期間圖像預處理的具體細節可在 此處 找到。圖像會被調整大小/重新縮放至相同的分辨率(224x224),並在 RGB 通道上進行歸一化,均值為 (0.5, 0.5, 0.5),標準差為 (0.5, 0.5, 0.5)。
預訓練
模型在 TPUv3 硬件(8 核)上進行訓練。所有模型變體均以 4096 的批量大小和 10k 步的學習率預熱進行訓練。對於 ImageNet,作者發現額外應用全局範數為 1 的梯度裁剪是有益的。訓練分辨率為 224。
評估結果
關於多個圖像分類基準的評估結果,請參考原論文的表 2 和表 5。需要注意的是,在微調時,使用更高的分辨率(384x384)可獲得最佳結果。當然,增加模型大小也會提高性能。
📄 許可證
本項目採用 Apache-2.0 許可證。
BibTeX 條目和引用信息
@misc{wu2020visual,
title={Visual Transformers: Token-based Image Representation and Processing for Computer Vision},
author={Bichen Wu and Chenfeng Xu and Xiaoliang Dai and Alvin Wan and Peizhao Zhang and Zhicheng Yan and Masayoshi Tomizuka and Joseph Gonzalez and Kurt Keutzer and Peter Vajda},
year={2020},
eprint={2006.03677},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
@inproceedings{deng2009imagenet,
title={Imagenet: A large-scale hierarchical image database},
author={Deng, Jia and Dong, Wei and Socher, Richard and Li, Li-Jia and Li, Kai and Fei-Fei, Li},
booktitle={2009 IEEE conference on computer vision and pattern recognition},
pages={248--255},
year={2009},
organization={Ieee}
}