timm 文档

从 Hugging Face Hub 分享和加载模型

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

从 Hugging Face Hub 分享和加载模型

timm 库内置了与 Hugging Face Hub 的集成,可以轻松地从 🤗 Hub 分享和加载模型。

在这篇简短的指南中,我们将看到如何:

  1. 在 Hub 上分享一个 timm 模型
  2. 如何从 Hub 加载该模型

身份验证

首先,你需要确保已经安装了 huggingface_hub 包。

pip install huggingface_hub

然后,你需要进行身份验证。你可以通过运行以下命令来完成:

huggingface-cli login

或者,如果你正在使用 notebook,可以使用 notebook_login 辅助函数。

>>> from huggingface_hub import notebook_login
>>> notebook_login()

分享模型

>>> import timm
>>> model = timm.create_model('resnet18', pretrained=True, num_classes=4)

这里通常是你训练或微调模型的地方。为了本教程的简洁,我们将跳过这一步。

假设我们现在已经微调好了模型。下一步就是将其推送到 Hub!我们可以使用 timm.models.hub.push_to_hf_hub 函数来完成。

>>> model_cfg = dict(label_names=['a', 'b', 'c', 'd'])
>>> timm.models.push_to_hf_hub(model, 'resnet18-random', model_config=model_cfg)

运行上述代码会将模型推送到 Hub 上的 <your-username>/resnet18-random。现在你可以与朋友分享这个模型,或者在自己的代码中使用它!

加载模型

从 Hub 加载模型非常简单,只需调用 timm.create_model 并将 pretrained 参数设置为你想要加载的模型名称即可。在这里,我们将使用 nateraw/resnet18-random,也就是我们刚刚推送到 Hub 的模型。

>>> model_reloaded = timm.create_model('hf_hub:nateraw/resnet18-random', pretrained=True)
< > 在 GitHub 上更新