将任何 ML 框架与 Hub 集成
Hugging Face Hub 使与社区共享和托管模型变得容易。它支持开源生态系统中的 数十个库。我们一直在努力扩展这种支持,以推动协作机器学习的发展。huggingface_hub
库在此过程中起着关键作用,它允许任何 Python 脚本轻松推送和加载文件。
有四种主要方法可以将库与 Hub 集成
- 推送至 Hub:实现一种方法将模型上传至 Hub。这包括模型权重,以及 模型卡 以及运行模型所需的任何其他相关信息或数据(例如,训练日志)。此方法通常称为
push_to_hub()
。 - 从 Hub 下载:实现一种方法从 Hub 加载模型。该方法应下载模型配置/权重并加载模型。此方法通常称为
from_pretrained
或load_from_hub()
。 - 推理 API:使用我们的服务器免费对您的库支持的模型进行推理。
- 小部件:在 Hub 上模型的登录页面上显示小部件。它允许用户从浏览器快速尝试模型。
在本指南中,我们将重点关注前两个主题。我们将介绍您可以用来集成库的两种主要方法,以及它们的优点和缺点。本指南最后将对所有内容进行总结,以帮助您在两者之间进行选择。请记住,这些只是指南,您可以根据自己的需求进行调整。
如果您有兴趣了解推理和小部件,可以遵循 本指南。在两种情况下,如果您正在将库与 Hub 集成,并且希望在 我们的文档 中列出,请与我们联系。
灵活的方法:助手
将库集成到 Hub 的第一个方法是实际自己实现 push_to_hub
和 from_pretrained
方法。这使您能够完全灵活地控制要上传/下载哪些文件以及如何处理特定于您框架的输入。您可以参考这两个 上传文件 和 下载文件 指南来了解如何执行此操作。例如,FastAI 集成就是如此实现的(参见 push_to_hub_fastai() 和 from_pretrained_fastai())。
库之间的实现可能有所不同,但工作流程通常类似。
from_pretrained
这通常是如何 from_pretrained
方法看起来的样子
def from_pretrained(model_id: str) -> MyModelClass:
# Download model from Hub
cached_model = hf_hub_download(
repo_id=repo_id,
filename="model.pkl",
library_name="fastai",
library_version=get_fastai_version(),
)
# Load model
return load_model(cached_model)
push_to_hub
push_to_hub
方法通常需要更多复杂性来处理仓库创建、生成模型卡和保存权重。一种常见的方法是将所有这些文件保存到一个临时文件夹中,上传它,然后删除它。
def push_to_hub(model: MyModelClass, repo_name: str) -> None:
api = HfApi()
# Create repo if not existing yet and get the associated repo_id
repo_id = api.create_repo(repo_name, exist_ok=True)
# Save all files in a temporary directory and push them in a single commit
with TemporaryDirectory() as tmpdir:
tmpdir = Path(tmpdir)
# Save weights
save_model(model, tmpdir / "model.safetensors")
# Generate model card
card = generate_model_card(model)
(tmpdir / "README.md").write_text(card)
# Save logs
# Save figures
# Save evaluation metrics
# ...
# Push to hub
return api.upload_folder(repo_id=repo_id, folder_path=tmpdir)
这当然只是一个例子。如果您对更复杂的处理(删除远程文件、动态上传权重、在本地持久化权重等)感兴趣,请参考 上传文件 指南。
局限性
虽然灵活,但这种方法也有一些缺点,特别是在维护方面。Hugging Face 用户在使用 huggingface_hub
时,通常会使用其他功能。例如,当从 Hub 加载文件时,通常会提供以下参数
token
:从私有仓库下载revision
:从特定分支下载cache_dir
:将文件缓存到特定目录force_download
/local_files_only
:是否重用缓存proxies
:配置 HTTP 会话
在推送模型时,支持类似的参数
commit_message
:自定义提交信息private
:如果缺少私有仓库,则创建一个create_pr
:创建 PR 而不是推送到main
branch
:推送到分支而不是main
分支allow_patterns
/ignore_patterns
:过滤要上传的文件token
- ...
所有这些参数都可以添加到我们上面看到的实现中,并传递给 huggingface_hub
方法。但是,如果参数发生更改或添加了新功能,您将需要更新您的包。支持这些参数也意味着您需要维护更多文档。要了解如何减轻这些局限性,让我们跳到下一节 **类继承**。
更复杂的方法:类继承
正如我们上面所见,库中需要包含两个主要方法来将其集成到 Hub:上传文件 (push_to_hub
) 和下载文件 (from_pretrained
)。您可以自己实现这些方法,但这会带来一些问题。为了解决这个问题,huggingface_hub
提供了一种使用类继承的工具。让我们看看它是如何工作的!
在许多情况下,库已经使用 Python 类实现了其模型。该类包含模型的属性以及加载、运行、训练和评估模型的方法。我们的方法是扩展此类,使用 mixin 包含上传和下载功能。一个 Mixin 是一个类,旨在通过多重继承将一组特定功能扩展到现有类。huggingface_hub
提供了自己的 mixin,即 ModelHubMixin。关键是了解它的行为以及如何自定义它。
ModelHubMixin 类实现了 3 个 公开 方法 (push_to_hub
、save_pretrained
和 from_pretrained
)。这些是您的用户将使用您的库加载/保存模型时调用的方法。 ModelHubMixin 还定义了 2 个 私有 方法 (_save_pretrained
和 _from_pretrained
)。这些是您必须实现的。因此,要集成您的库,您应该
- 使您的模型类继承自 ModelHubMixin。
- 实现私有方法
- _save_pretrained():该方法以目录路径作为输入,并将模型保存到该路径。您必须在此方法中编写所有逻辑来转储您的模型:模型卡、模型权重、配置文件、训练日志和图形。此模型的任何相关信息都必须由此方法处理。 模型卡 特别重要,用于描述您的模型。查看 我们的实现指南 以获取更多详细信息。
- _from_pretrained():**类方法**,以
model_id
作为输入,返回一个已实例化的模型。该方法必须下载相关文件并加载它们。
- 您完成了!
使用 ModelHubMixin 的优势在于,一旦您处理完文件的序列化/加载,您就可以开始使用。您不必担心诸如仓库创建、提交、PR 或修订版之类的事情。 ModelHubMixin 还会确保公共方法得到文档化和类型标注,并且您将能够在 Hub 上查看模型的下载次数。所有这一切都由 ModelHubMixin 处理,并可供您的用户使用。
一个具体示例:PyTorch
上面我们所见的一个很好的例子是 PyTorchModelHubMixin,我们针对 PyTorch 框架的集成。这是一个现成的集成。
如何使用它?
以下是如何让任何用户从 Hub 加载/保存 PyTorch 模型
>>> import torch
>>> import torch.nn as nn
>>> from huggingface_hub import PyTorchModelHubMixin
# Define your Pytorch model exactly the same way you are used to
>>> class MyModel(
... nn.Module,
... PyTorchModelHubMixin, # multiple inheritance
... library_name="keras-nlp",
... tags=["keras"],
... repo_url="https://github.com/keras-team/keras-nlp",
... docs_url="https://keras.org.cn/keras_nlp/",
... # ^ optional metadata to generate model card
... ):
... def __init__(self, hidden_size: int = 512, vocab_size: int = 30000, output_size: int = 4):
... super().__init__()
... self.param = nn.Parameter(torch.rand(hidden_size, vocab_size))
... self.linear = nn.Linear(output_size, vocab_size)
... def forward(self, x):
... return self.linear(x + self.param)
# 1. Create model
>>> model = MyModel(hidden_size=128)
# Config is automatically created based on input + default values
>>> model.param.shape[0]
128
# 2. (optional) Save model to local directory
>>> model.save_pretrained("path/to/my-awesome-model")
# 3. Push model weights to the Hub
>>> model.push_to_hub("my-awesome-model")
# 4. Initialize model from the Hub => config has been preserved
>>> model = MyModel.from_pretrained("username/my-awesome-model")
>>> model.param.shape[0]
128
# Model card has been correctly populated
>>> from huggingface_hub import ModelCard
>>> card = ModelCard.load("username/my-awesome-model")
>>> card.data.tags
["keras", "pytorch_model_hub_mixin", "model_hub_mixin"]
>>> card.data.library_name
"keras-nlp"
实现
实现其实非常简单,完整的实现可以在这里找到 here.
- 首先,让你的类继承自
ModelHubMixin
from huggingface_hub import ModelHubMixin
class PyTorchModelHubMixin(ModelHubMixin):
(...)
- 实现
_save_pretrained
方法
from huggingface_hub import ModelHubMixin
class PyTorchModelHubMixin(ModelHubMixin):
(...)
def _save_pretrained(self, save_directory: Path) -> None:
"""Save weights from a Pytorch model to a local directory."""
save_model_as_safetensor(self.module, str(save_directory / SAFETENSORS_SINGLE_FILE))
- 实现
_from_pretrained
方法
class PyTorchModelHubMixin(ModelHubMixin):
(...)
@classmethod # Must be a classmethod!
def _from_pretrained(
cls,
*,
model_id: str,
revision: str,
cache_dir: str,
force_download: bool,
proxies: Optional[Dict],
resume_download: bool,
local_files_only: bool,
token: Union[str, bool, None],
map_location: str = "cpu", # additional argument
strict: bool = False, # additional argument
**model_kwargs,
):
"""Load Pytorch pretrained weights and return the loaded model."""
model = cls(**model_kwargs)
if os.path.isdir(model_id):
print("Loading weights from local directory")
model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
return cls._load_as_safetensor(model, model_file, map_location, strict)
model_file = hf_hub_download(
repo_id=model_id,
filename=SAFETENSORS_SINGLE_FILE,
revision=revision,
cache_dir=cache_dir,
force_download=force_download,
proxies=proxies,
resume_download=resume_download,
token=token,
local_files_only=local_files_only,
)
return cls._load_as_safetensor(model, model_file, map_location, strict)
就是这样!你的库现在可以让用户上传和下载到 Hub 上的文件。
高级用法
在上一节中,我们简单介绍了 ModelHubMixin 的工作原理。在本节中,我们将了解一些更高级的功能,以改进你的库与 Hugging Face Hub 的集成。
模型卡
ModelHubMixin 会为你生成模型卡。模型卡是与模型一起提供的文件,提供有关模型的重要信息。在底层,模型卡是带有附加元数据的简单 Markdown 文件。模型卡对于可发现性、可重复性和共享至关重要!查看 模型卡指南以获取更多详细信息。
半自动生成模型卡是确保使用你的库推送的所有模型都会共享通用元数据的好方法:library_name
、tags
、license
、pipeline_tag
等。这使得由你的库支持的所有模型都可以在 Hub 上轻松搜索,并为在 Hub 上的着陆用户提供一些资源链接。你可以直接在继承自 ModelHubMixin 时定义元数据。
class UniDepthV1(
nn.Module,
PyTorchModelHubMixin,
library_name="unidepth",
repo_url="https://github.com/lpiccinelli-eth/UniDepth",
docs_url=...,
pipeline_tag="depth-estimation",
license="cc-by-nc-4.0",
tags=["monocular-metric-depth-estimation", "arxiv:1234.56789"]
):
...
默认情况下,将使用你提供的的信息生成通用的模型卡(例如:pyp1/VoiceCraft_giga830M)。但你也可以定义自己的模型卡模板!
在这个例子中,使用 VoiceCraft
类推送的所有模型都将自动包含引用部分和许可证详细信息。有关如何定义模型卡模板的更多详细信息,请查看 模型卡指南.
MODEL_CARD_TEMPLATE = """
---
# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
# Doc / guide: https://huggingface.co/docs/hub/model-cards
{{ card_data }}
---
This is a VoiceCraft model. For more details, please check out the official Github repo: https://github.com/jasonppy/VoiceCraft. This model is shared under a Attribution-NonCommercial-ShareAlike 4.0 International license.
## Citation
@article{peng2024voicecraft,
author = {Peng, Puyuan and Huang, Po-Yao and Li, Daniel and Mohamed, Abdelrahman and Harwath, David},
title = {VoiceCraft: Zero-Shot Speech Editing and Text-to-Speech in the Wild},
journal = {arXiv},
year = {2024},
}
"""
class VoiceCraft(
nn.Module,
PyTorchModelHubMixin,
library_name="voicecraft",
model_card_template=MODEL_CARD_TEMPLATE,
...
):
...
最后,如果你想使用动态值扩展模型卡生成过程,可以覆盖 generate_model_card()
方法。
from huggingface_hub import ModelCard, PyTorchModelHubMixin
class UniDepthV1(nn.Module, PyTorchModelHubMixin, ...):
(...)
def generate_model_card(self, *args, **kwargs) -> ModelCard:
card = super().generate_model_card(*args, **kwargs)
card.data.metrics = ... # add metrics to the metadata
card.text += ... # append section to the modelcard
return card
配置
ModelHubMixin 会为你处理模型配置。它会自动检查你在实例化模型时输入的值,并将它们序列化到 config.json
文件中。这提供了 2 个好处
- 用户将能够使用与你完全相同的参数重新加载模型。
- 拥有
config.json
文件会自动在 Hub 上启用分析(即“下载”计数)。
但它在实践中是如何工作的?一些规则使得从用户的角度来看,这个过程尽可能地顺利。
- 如果你的
__init__
方法期望一个config
输入,它将自动保存在仓库中,文件名是config.json
。 - 如果
config
输入参数用数据类类型注释(例如:config: Optional[MyConfigClass] = None
),那么config
值将为你正确反序列化。 - 在初始化时传递的所有值也将存储在配置文件中。这意味着你并不一定需要期望一个
config
输入才能从中受益。
示例
class MyModel(ModelHubMixin):
def __init__(value: str, size: int = 3):
self.value = value
self.size = size
(...) # implement _save_pretrained / _from_pretrained
model = MyModel(value="my_value")
model.save_pretrained(...)
# config.json contains passed and default values
{"value": "my_value", "size": 3}
但是,如果一个值不能作为 JSON 序列化?默认情况下,该值将在保存配置文件时被忽略。但是,在某些情况下,你的库已经期望一个不能序列化的自定义对象作为输入,并且你不想更新你的内部逻辑来更新它的类型。别担心!在继承自 ModelHubMixin 时,你可以为任何类型传递自定义编码器/解码器。这需要做更多工作,但确保你的内部逻辑在将你的库与 Hub 集成时不会受到影响。
这是一个具体的示例,其中一个类期望一个 argparse.Namespace
配置作为输入
class VoiceCraft(nn.Module):
def __init__(self, args):
self.pattern = self.args.pattern
self.hidden_size = self.args.hidden_size
...
一个解决方案可能是将 __init__
签名更新为 def __init__(self, pattern: str, hidden_size: int)
并更新所有实例化你的类的代码片段。这是一种完全有效的解决方法,但它可能会破坏使用你的库的下游应用程序。
另一个解决方案是提供一个简单的编码器/解码器,将 argparse.Namespace
转换为字典。
from argparse import Namespace
class VoiceCraft(
nn.Module,
PyTorchModelHubMixin, # inherit from mixin
coders={
Namespace : (
lambda x: vars(x), # Encoder: how to convert a `Namespace` to a valid jsonable value?
lambda data: Namespace(**data), # Decoder: how to reconstruct a `Namespace` from a dictionary?
)
}
):
def __init__(self, args: Namespace): # annotate `args`
self.pattern = self.args.pattern
self.hidden_size = self.args.hidden_size
...
在上面的代码片段中,类的内部逻辑和 __init__
签名都没有改变。这意味着你的库的所有现有代码片段都将继续工作。为了实现这一点,我们必须
- 继承自混入(在本例中是
PytorchModelHubMixin
)。 - 在继承时传递一个
coders
参数。这是一个字典,其中键是你想要处理的自定义类型。值是元组(encoder, decoder)
。- 编码器期望一个指定类型的对象作为输入,并返回一个可 JSON 化的值。这将在使用
save_pretrained
保存模型时使用。 - 解码器期望原始数据(通常是字典)作为输入,并重建初始对象。这将在使用
from_pretrained
加载模型时使用。
- 编码器期望一个指定类型的对象作为输入,并返回一个可 JSON 化的值。这将在使用
- 在
__init__
签名中添加类型注释。这对于让混入知道类期望哪种类型以及因此要使用哪个解码器非常重要。
为了简单起见,上面的示例中的编码器/解码器函数并不健壮。对于具体的实现,你很可能需要适当地处理角点情况。
快速比较
让我们快速总结一下我们看到的两种方法及其优缺点。下面的表格仅供参考。你的框架可能有一些你需要解决的特殊情况。本指南仅用于提供指导和想法,说明如何处理集成。无论如何,如果你有任何问题,请随时联系我们!
集成 | 使用助手 | 使用 ModelHubMixin |
---|---|---|
用户体验 | model = load_from_hub(...) push_to_hub(model, ...) | model = MyModel.from_pretrained(...) model.push_to_hub(...) |
灵活性 | 非常灵活。 你可以完全控制实现。 | 灵活性较低。 你的框架必须有一个模型类。 |
维护 | 需要更多维护来添加配置支持和新功能。可能还需要修复用户报告的错误。 | 维护量较少,因为与 Hub 的大多数交互都已在 huggingface_hub 中实现。 |
文档/类型注释 | 需要手动编写。 | 部分由 huggingface_hub 处理。 |
下载计数器 | 需要手动处理。 | 如果类具有 config 属性,则默认启用。 |
模型卡 | 需要手动处理 | 默认情况下会生成,其中包含 library_name、tags 等。 |