加载模型

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 上更新