載入模型

Transformers 提供了許多預訓練模型，只需一行程式碼即可使用。它需要一個模型類和 from_pretrained() 方法。

呼叫 from_pretrained() 以下載並載入儲存在 Hugging Face Hub 上的模型權重和配置。

如果可用，from_pretrained() 方法會載入以 safetensors 檔案格式儲存的權重。傳統上，PyTorch 模型權重使用 pickle 工具進行序列化，而該工具被認為是不安全的。Safetensor 檔案更安全且載入速度更快。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype="auto", device_map="auto")

本指南解釋了模型是如何載入的，載入模型的不同方式，如何解決超大模型的記憶體問題，以及如何載入自定義模型。

模型和配置

所有模型都有一個 `configuration.py` 檔案，其中包含隱藏層數量、詞彙大小、啟用函式等特定屬性。您還會找到一個 `modeling.py` 檔案，該檔案定義了每個層中的層和數學運算。`modeling.py` 檔案獲取 `configuration.py` 中的模型屬性並相應地構建模型。此時，您有一個具有隨機權重的模型，需要對其進行訓練才能輸出有意義的結果。

架構指的是模型的骨架，而檢查點 指的是給定架構的模型權重。例如，BERT 是一種架構，而 google-bert/bert-base-uncased 是一個檢查點。您會看到模型這個術語與架構和檢查點互換使用。

您可以載入兩種一般型別的模型：

一個基礎模型，例如 AutoModel 或 LlamaModel，它輸出隱藏狀態。
一個帶有特定頭部的模型，例如 AutoModelForCausalLM 或 LlamaForCausalLM，用於執行特定任務。

對於每種模型型別，每個機器學習框架（PyTorch、TensorFlow、Flax）都有一個單獨的類。選擇您正在使用的框架的相應字首。

PyTorch

TensorFlow

Flax

模型類

要獲取預訓練模型，您需要將權重載入到模型中。這透過呼叫 from_pretrained() 完成，該方法接受來自 Hugging Face Hub 或本地目錄的權重。

有兩種模型類：AutoModel 類和模型特定類。

自動模型

模型特定類

大型模型

大型預訓練模型需要大量記憶體才能載入。載入過程包括：

建立一個帶有隨機權重的模型
載入預訓練權重
將預訓練權重放置在模型上

您需要足夠的記憶體來容納模型權重的兩個副本（隨機和預訓練），這可能取決於您的硬體。在分散式訓練環境中，這更具挑戰性，因為每個程序都會載入一個預訓練模型。

Transformers 透過快速初始化、分片檢查點、Accelerate 的大型模型推理功能以及支援較低位資料型別來減少這些記憶體相關的挑戰。

分片檢查點

save_pretrained() 方法會自動對大於 10GB 的檢查點進行分片。

每個分片會在前一個分片載入後按順序載入，將記憶體使用限制在僅模型大小和最大分片大小。

預設情況下，`max_shard_size` 引數設定為每個分片 5GB，因為這樣更容易在免費 GPU 例項上執行而不會耗盡記憶體。

例如，在 save_pretrained() 中為 BioMistral/BioMistral-7B 建立一些分片檢查點。

from transformers import AutoModel
import tempfile
import os

model = AutoModel.from_pretrained("biomistral/biomistral-7b")
with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="5GB")
    print(sorted(os.listdir(tmp_dir)))

使用 from_pretrained() 重新載入分片檢查點。

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir)
    new_model = AutoModel.from_pretrained(tmp_dir)

分片檢查點也可以直接用 load_sharded_checkpoint() 載入。

from transformers.modeling_utils import load_sharded_checkpoint

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="5GB")
    load_sharded_checkpoint(model, tmp_dir)

save_pretrained() 方法會建立一個索引檔案，將引數名對映到它們儲存的檔案。索引檔案有兩個鍵，`metadata` 和 `weight_map`。

import json

with tempfile.TemporaryDirectory() as tmp_dir:
    model.save_pretrained(tmp_dir, max_shard_size="5GB")
    with open(os.path.join(tmp_dir, "model.safetensors.index.json"), "r") as f:
        index = json.load(f)

print(index.keys())

`metadata` 鍵提供了模型總大小。

index["metadata"]
{'total_size': 28966928384}

`weight_map` 鍵將每個引數對映到它儲存的分片。

index["weight_map"]
{'lm_head.weight': 'model-00006-of-00006.safetensors',
 'model.embed_tokens.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.input_layernorm.weight': 'model-00001-of-00006.safetensors',
 'model.layers.0.mlp.down_proj.weight': 'model-00001-of-00006.safetensors',
 ...
}

大型模型推理

請確保您已安裝 Accelerate v0.9.0 和 PyTorch v1.9.0 或更高版本以使用此功能！

from_pretrained() 得到了 Accelerate 的大型模型推理功能的增強。

大型模型推理在 PyTorch 元裝置上建立一個模型骨架。元裝置不儲存任何實際資料，只儲存元資料。

隨機初始化的權重只在載入預訓練權重時建立，以避免同時在記憶體中維護模型的兩個副本。最大記憶體使用量僅為模型的大小。

在設計裝置對映中瞭解更多關於裝置放置的資訊。

大型模型推理的第二個功能與模型骨架中權重的載入和分配方式有關。模型權重被分配到所有可用裝置，從最快的裝置（通常是 GPU）開始，然後將任何剩餘的權重解除安裝到較慢的裝置（CPU 和硬碟）。

這兩個功能結合起來減少了大型預訓練模型的記憶體使用和載入時間。

將 device_map 設定為 ` "auto" ` 以啟用大型模型推理。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device_map="auto")

您還可以手動將層分配到 `device_map` 中的裝置。它應該將所有模型引數對映到一個裝置，但如果整個層都在同一個裝置上，則無需詳細說明層的所有子模組去向何處。

訪問 `hf_device_map` 屬性以檢視模型如何在裝置之間分佈。

device_map = {"model.layers.1": 0, "model.layers.14": 1, "model.layers.31": "cpu", "lm_head": "disk"}
model.hf_device_map

模型資料型別

PyTorch 模型權重預設初始化為 `torch.float32`。以不同資料型別（例如 `torch.float16`）載入模型需要額外的記憶體，因為模型會以所需的資料型別再次載入。

顯式設定 torch_dtype 引數，以直接以所需的資料型別初始化模型，而不是兩次載入權重（`torch.float32` 然後是 `torch.float16`）。您也可以設定 `torch_dtype="auto"` 以自動載入權重，使其資料型別與儲存的資料型別相同。

特定 dtype

自動 dtype

`torch_dtype` 引數也可以在 AutoConfig 中配置，用於從頭開始例項化的模型。

import torch
from transformers import AutoConfig, AutoModel

my_config = AutoConfig.from_pretrained("google/gemma-2b", torch_dtype=torch.float16)
model = AutoModel.from_config(my_config)

自定義模型

自定義模型建立在 Transformers 的配置和建模類之上，支援 AutoClass API，並使用 from_pretrained() 載入。不同之處在於建模程式碼不來自 Transformers。

載入自定義模型時請格外小心。雖然 Hub 包含對每個儲存庫的惡意軟體掃描，但您仍應小心，以避免無意中執行惡意程式碼。

在 from_pretrained() 中設定 `trust_remote_code=True` 以載入自定義模型。

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained("sgugger/custom-resnet50d", trust_remote_code=True)

作為額外的安全層，從特定修訂版載入自定義模型，以避免載入可能已更改的模型程式碼。提交雜湊可以從模型提交歷史記錄中複製。

commit_hash = "ed94a7c6247d8aedce4647f00f20de6875b5b292"
model = AutoModelForImageClassification.from_pretrained(
    "sgugger/custom-resnet50d", trust_remote_code=True, revision=commit_hash
)

有關更多資訊，請參閱自定義模型指南。

< > 在 GitHub 上更新