PyTorchModelHubMixin：弥合 Hugging Face 上自定义 AI 模型之间的鸿沟

简介

Hugging Face 提供了一系列实用功能，包括免费 AI 模型托管、版本控制、每月下载量指标等。

对于使用自定义 AI 模型架构的用户来说，访问这些功能可能会遇到显著障碍。而 **PyTorchModelHubMixin** 正是为此而生，它能无缝实现这些模型的保存、加载和推送到 Hub。

在这篇博文中，我们将以 pxia 为参考，探讨如何利用 PyTorchModelHubMixin 构建您自己的 Hugging Face 集成库。

PyTorchModelHubMixin：主要特性与功能

虽然大多数著名的 AI 模型已经成为 transformers 库的一部分，但一些新兴的模型正在使用 mixin 类，其中包括 lerobot、ibm/biomed.sm.mv-te-84m、Salesforce/moirai-1.1-R-large、nvidia/domain-classifier、ZhengPeng7/BiRefNet 等。

尽管 Transformers 库提供了广泛的功能套件，但 PyTorchModelHubMixin 在灵活性方面提供了更为精简的功能集。为了进行详细比较，我们使用 pxia 作为参考，编制了以下综合表格，展示了使用 PyTorchModelHubMixin 可以实现的功能。

• ✅ 库中已包含
• 🟡 库中未包含，但可添加该功能
• ❌ 无法添加到库中

用法

PyTorchModelHubMixin 的用法非常简单，您只需将其集成到模型的类继承中即可。请看以下极简示例作为参考：

from torch import nn
+from huggingface_hub import PyTorchModelHubMixin

class ANN(nn.Module,
+          PyTorchModelHubMixin):
  def __init__(self,a,b):
    super().__init__()
    self.layer = nn.Linear(a,b, bias=False)
  def forward(self,inputs):
    return self.layer(inputs)

您所需要做的就这些 :)

如果您想在 readme 文件中添加一些额外的元数据，可以按照文档中的定义，将它们添加到继承中。

class MyModel(
        nn.Module,
        PyTorchModelHubMixin,
        library_name="keras-nlp",
        repo_url="https://github.com/keras-team/keras-nlp",
        docs_url="https://keras.org.cn/keras_nlp/",
        tags = ["demo","architecture","tensorflow"]
        # ^ optional metadata to generate model card
    ):

模型的推送和加载

通过继承 `PyTorchModelHubMixin`，您的模型将继承三个便捷的方法：`save_pretrained`、`from_pretrained` 和 `push_to_hub`。这些方法的功能类似于 Transformers 库中的同名方法，可实现模型的无缝共享和加载。让我们深入了解一个实际示例，说明如何在实际场景中使用 mixin。

1- 定义您的模型架构

from torch import nn
from huggingface_hub import PyTorchModelHubMixin

class ANN(nn.Module,
          PyTorchModelHubMixin):
  def __init__(self,a,b):
    super().__init__()
    self.layer = nn.Linear(a,b, bias=False)
  def forward(self,inputs):
    return self.layer(inputs)

2- 初始化并将您的模型推送到 Hub，您也可以使用 `save_pretrained` 在本地保存模型。

model = ANN(1,2)
model.push_to_hub("repo_id",token=TOKEN)

3- 在此处加载模型，您可以使用 **ModelClass** 直接加载模型，无需再次重新初始化模型。

# No init parameters, no manual initialization or trying to remember the __init__ parameters 
HF_model = ANN.from_pretrained("repo_id_or_path") # you can also pass your revision parameter here to load a specific sha,branch,pull of the model

`from_pretrained` 在模型本地保存时也同样适用。

您可能也注意到，在步骤 (3) 中我们没有传递任何初始化参数，这是因为我们所有的参数都在步骤 (2) 中捕获并保存在 `config.json` 文件中。简而言之，您只需定义一次初始化参数即可。

边缘情况 🛠️

在本节中，我们将概述开发 Hugging Face 集成 AI 模型的基本最佳实践和注意事项。

为确保您的模型按预期运行并顺利集成，请牢记以下两项关键建议：

• 避免在 `__init__` 方法中加载任何本地文件

import torch
from torch import nn
from huggingface_hub import PyTorchModelHubMixin
weights="folder/weights.pt"

class ANN(nn.Module,
          PyTorchModelHubMixin):
  def __init__(self,a,b):
    super().__init__()
    self.layer = nn.Linear(a,b, bias=False)
-    self.load_state_dict(torch.load(weights))

  def forward(self,inputs):
    return self.layer(inputs)

model = ANN(1,2)
+model.load_state_dict(torch.load(weights))

• 如果您的输入类型不可序列化，请定义 `coders` 参数

可序列化的 `__init__` 参数可以在此处找到

如果您的输入参数不受 mixin 支持，您可以在 `coders` 参数中添加序列化和反序列化逻辑（这将定义我们如何将这些参数存储在 `config.json` 文件中）

import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
from omegaconf import OmegaConf
from omegaconf.dictconfig import DictConfig

default_conf = OmegaConf.create({"a": 2, "b": 1})

# converts parameter to python dict
+def serialize(x):
+    return OmegaConf.to_container(x)

# when loading the model we convert the init parameter back to its original type
+def deserialize(x):
+    return OmegaConf.create(x)


class ANN(
    nn.Module,
    PyTorchModelHubMixin,
+    coders={
+        DictConfig: (
+            lambda x: serialize(x),
+            lambda data: deserialize(data),
+        )
+    },
):
    # infront of the parameter you need to define the input type beforehand so we can figure out
    # which serialization logic goes to which parameter
    def __init__(self,
-                     cfg = default_conf):
+                     cfg:DictConfig = default_conf): 
        super().__init__()
        self.layer = nn.Linear(cfg.a, cfg.b, bias=False)

    def forward(self, inputs):
        return self.layer(inputs)

未来工作

如需全面了解本博文中提到的功能，请参阅 pxia。Pxia 是一个 AI 库，旨在作为一个模板，通过实际示例帮助用户开发自己的项目。

如果您想为 PEFT、量化或您认为对用户有价值的任何其他功能做出贡献，请提交拉取请求或在 GitHub 仓库中创建问题。

如果您觉得这篇博文有帮助，请考虑点赞 ❤️ヾ(≧▽≦*)o