ONNX 模型中心

ONNX 模型中心是從 ONNX 模型動物園開始使用最先進的預先訓練 ONNX 模型的一個簡單快速的方法。此外,這也讓研究人員和模型開發人員有機會與更廣泛的社群分享他們預先訓練的模型。

安裝

ONNX 模型中心在 ONNX 1.11.0 之後可用。

基本用法

ONNX 模型中心能夠從任何 git 儲存庫下載、列出和查詢訓練好的模型,並且預設為官方的 ONNX 模型動物園。在本節中,我們示範一些基本功能。

首先,請使用以下方式導入中心

from onnx import hub

依名稱下載模型:

load 函式預設會在模型動物園中搜尋具有相符名稱的最新模型,將此模型下載到本機快取,並將模型載入到 ModelProto 物件中,以供 ONNX 執行階段使用。

model = hub.load("resnet50")

從自訂儲存庫下載:

任何具有適當結構的儲存庫都可以作為 ONNX 模型中心。若要從其他中心下載,或指定主要模型中心上的特定分支或提交,可以提供 repo 參數

model = hub.load("resnet50", repo="onnx/models:771185265efbdc049fb223bd68ab1aeb1aecde76")

列出和檢查模型:

模型中心提供 API 來查詢模型動物園,以了解有關可用模型的更多資訊。這不會下載模型,而是只會傳回有關符合給定引數的模型資訊

# List all models in the onnx/models:main repo
all_models = hub.list_models()

# List all versions/opsets of a specific model
mnist_models = hub.list_models(model="mnist")

# List all models matching a given "tag"
vision_models = hub.list_models(tags=["vision"])

也可以使用 get_model_info 函式,在下載之前檢查模型的中繼資料

print(hub.get_model_info(model="mnist", opset=8))

這將會列印類似以下的內容

ModelInfo(
    model=MNIST,
    opset=8,
    path=vision/classification/mnist/model/mnist-8.onnx,
    metadata={
     'model_sha': '2f06e72de813a8635c9bc0397ac447a601bdbfa7df4bebc278723b958831c9bf',
     'model_bytes': 26454,
     'tags': ['vision', 'classification', 'mnist'],
     'io_ports': {
        'inputs': [{'name': 'Input3', 'shape': [1, 1, 28, 28], 'type': 'tensor(float)'}],
        'outputs': [{'name': 'Plus214_Output_0', 'shape': [1, 10], 'type': 'tensor(float)'}]},
     'model_with_data_path': 'vision/classification/mnist/model/mnist-8.tar.gz',
     'model_with_data_sha': '1dd098b0fe8bc750585eefc02013c37be1a1cae2bdba0191ccdb8e8518b3a882',
     'model_with_data_bytes': 25962}
)

本機快取

ONNX 模型中心會在本機快取已下載的模型 (位於可設定的位置),以便後續對 hub.load 的呼叫不需要網路連線。

預設快取位置

中心用戶端會依此順序尋找以下預設快取位置

  1. 如果已定義 ONNX_HOME 環境變數,則為 $ONNX_HOME/hub

  2. 如果已定義 XDG_CACHE_HOME 環境變數,則為 $XDG_CACHE_HOME/hub

  3. ~/.cache/onnx/hub,其中 ~ 是使用者主目錄

設定快取位置

若要手動設定快取位置,請使用

hub.set_dir("my/cache/directory")

此外,可以使用以下方式檢查快取位置

print(hub.get_dir())

其他快取詳細資料

若要清除模型快取,可以使用類似 shutilos 的 Python 公用程式來刪除快取目錄。此外,可以使用 force_reload 選項來選擇覆寫快取模型

model = hub.load("resnet50", force_reload=True)

我們包含此旗標是為了完整性,但請注意,快取中的模型會使用 sha256 雜湊值來消除歧義,因此在正常使用情況下不需要 force_reload 旗標。最後,我們注意到模型快取目錄結構會鏡射資訊清單的 model_path 欄位所指定的目錄結構,但是檔案名稱會使用模型 SHA256 雜湊值消除歧義。

如此一來,模型快取對人類來說是可讀的,可以在多個模型版本之間消除歧義,並且可以在不同中心之間重複使用快取模型 (如果它們具有相同的名稱和雜湊值)。

架構

ONNX Hub Architecture

ONNX 中心包含兩個主要元件,用戶端和伺服器。用戶端程式碼目前包含在 onnx 套件中,並且可以指向伺服器,形式為 github 儲存庫中託管的 ONNX_HUB_MANIFEST.json,例如 ONNX 模型動物園中的那個。此資訊清單檔案是一個 JSON 文件,其中列出所有模型及其元數據,並且設計為與程式設計語言無關。以下是一個結構良好的模型資訊清單項目的範例

{
 "model": "BERT-Squad",
 "model_path": "text/machine_comprehension/bert-squad/model/bertsquad-8.onnx",
 "onnx_version": "1.3",
 "opset_version": 8,
 "metadata": {
     "model_sha": "cad65b9807a5e0393e4f84331f9a0c5c844d9cc736e39781a80f9c48ca39447c",
     "model_bytes": 435882893,
     "tags": ["text", "machine comprehension", "bert-squad"],
     "io_ports": {
         "inputs": [
             {
                 "name": "unique_ids_raw_output___9:0",
                 "shape": ["unk__475"],
                 "type": "tensor(int64)"
             },
             {
                 "name": "segment_ids:0",
                 "shape": ["unk__476", 256],
                 "type": "tensor(int64)"
             },
             {
                 "name": "input_mask:0",
                 "shape": ["unk__477", 256],
                 "type": "tensor(int64)"
             },
             {
                 "name": "input_ids:0",
                 "shape": ["unk__478", 256],
                 "type": "tensor(int64)"
             }
         ],
         "outputs": [
             {
                 "name": "unstack:1",
                 "shape": ["unk__479", 256],
                 "type": "tensor(float)"
             },
             {
                 "name": "unstack:0",
                 "shape": ["unk__480", 256],
                 "type": "tensor(float)"
             },
             {
                 "name": "unique_ids:0",
                 "shape": ["unk__481"],
                 "type": "tensor(int64)"
             }
         ]
     },
     "model_with_data_path": "text/machine_comprehension/bert-squad/model/bertsquad-8.tar.gz",
     "model_with_data_sha": "c8c6c7e0ab9e1333b86e8415a9d990b2570f9374f80be1c1cb72f182d266f666",
     "model_with_data_bytes": 403400046
 }
}

這些重要的欄位是

  • model:用於查詢的模型名稱

  • model_path:儲存在 Git LFS 中的模型相對路徑。

  • onnx_version:模型的 ONNX 版本

  • opset_version:opset 的版本。如果未指定,則用戶端會下載最新的 opset。

  • metadata/model_sha:選用的模型 sha 規格,用於提高下載安全性

  • metadata/tags:選用的高階標籤,可協助使用者依給定類型尋找模型

metadata 欄位中的所有其他欄位對於用戶端而言都是選用的,但會為使用者提供重要的詳細資料。

新增至 ONNX 模型中心

貢獻官方模型

將模型新增至官方 onnx/models 版本模型中心的最簡單方法是遵循這些指南 來貢獻您的模型。一旦貢獻,請確保您的模型在其 README.md 中具有 Markdown 表格 (範例)。模型中心資訊清單產生器會從這些 Markdown 表格中提取資訊。若要執行產生器

git clone https://github.com/onnx/models.git
git lfs pull --include="*" --exclude=""
cd models/workflow_scripts
python generate_onnx_hub_manifest.py

一旦產生新的資訊清單,請新增並將其提交至 onnx/models 的提取要求中

託管您自己的 ONNX 模型中心

若要託管您自己的模型中心,請將 ONNX_HUB_MANIFEST.json 新增至您的 github 儲存庫的頂層 (範例)。您的資訊清單項目至少應包含本文檔 架構章節中提及的欄位。提交後,請檢查是否可以使用此文檔的「從自訂儲存庫下載」章節下載模型。

如有任何問題,請提出