torch.hub¶
Pytorch Hub 是一个预训练模型存储库,旨在促进研究的可重复性。
发布模型¶
Pytorch Hub 支持通过添加简单的 hubconf.py
文件将预训练模型(模型定义和预训练权重)发布到 GitHub 存储库;
hubconf.py
可以有多个入口点。每个入口点都定义为一个 Python 函数(例如:您要发布的预训练模型)。
def entrypoint_name(*args, **kwargs):
# args & kwargs are optional, for models which take positional/keyword arguments.
...
如何实现入口点?¶
以下代码片段指定了 resnet18
模型的入口点,如果我们在 pytorch/vision/hubconf.py
中扩展实现。在大多数情况下,在 hubconf.py
中导入正确的函数就足够了。这里我们只是想使用扩展版本作为示例来展示其工作原理。您可以在 pytorch/vision 存储库 中查看完整脚本
dependencies = ['torch']
from torchvision.models.resnet import resnet18 as _resnet18
# resnet18 is the name of entrypoint
def resnet18(pretrained=False, **kwargs):
""" # This docstring shows up in hub.help()
Resnet18 model
pretrained (bool): kwargs, load pretrained weights into the model
"""
# Call the model, load pretrained weights
model = _resnet18(pretrained=pretrained, **kwargs)
return model
dependencies
变量是加载模型所需的包名称的**列表**。请注意,这可能与训练模型所需的依赖项略有不同。args
和kwargs
将传递给真正的可调用函数。函数的文档字符串用作帮助消息。它解释了模型的功能以及允许的位置/关键字参数。强烈建议在此处添加一些示例。
入口点函数可以返回模型 (nn.module),也可以返回辅助工具以使用户工作流程更流畅,例如标记器。
以下划线为前缀的可调用对象被视为辅助函数,不会显示在
torch.hub.list()
中。预训练权重可以存储在 GitHub 存储库中的本地位置,也可以通过
torch.hub.load_state_dict_from_url()
加载。如果小于 2GB,建议将其附加到 项目发布 并使用发布中的 URL。在上面的示例中,torchvision.models.resnet.resnet18
处理pretrained
,或者您也可以在入口点定义中放置以下逻辑。
if pretrained:
# For checkpoint saved in local GitHub repo, e.g. <RELATIVE_PATH_TO_CHECKPOINT>=weights/save.pth
dirname = os.path.dirname(__file__)
checkpoint = os.path.join(dirname, <RELATIVE_PATH_TO_CHECKPOINT>)
state_dict = torch.load(checkpoint)
model.load_state_dict(state_dict)
# For checkpoint saved elsewhere
checkpoint = 'https://download.pytorch.org/models/resnet18-5c106cde.pth'
model.load_state_dict(torch.hub.load_state_dict_from_url(checkpoint, progress=False))
重要通知¶
发布的模型至少应位于分支/标签中。它不能是随机提交。
从 Hub 加载模型¶
Pytorch Hub 提供了方便的 API,可以通过 torch.hub.list()
探索 Hub 中所有可用的模型,通过 torch.hub.help()
显示文档字符串和示例,并使用 torch.hub.load()
加载预训练模型。
- torch.hub.list(github, force_reload=False, skip_validation=False, trust_repo=None, verbose=True)[源代码]¶
列出
github
指定的存储库中可用的所有可调用入口点。- 参数
github (str) – 格式为“repo_owner/repo_name[:ref]”的字符串,其中包含可选的 ref(标签或分支)。如果未指定
ref
,则假设默认分支为main
(如果存在),否则为master
。示例:'pytorch/vision:0.10'force_reload (bool,可选) – 是否丢弃现有缓存并强制重新下载。默认为
False
。skip_validation (bool,可选) – 如果为
False
,则 torchhub 将检查由github
参数指定的 branch 或 commit 是否正确属于 repo 所有者。这将向 GitHub API 发出请求;可以通过设置GITHUB_TOKEN
环境变量来指定非默认的 GitHub token。默认为False
。trust_repo (bool,str 或 None) –
"check"
、True
、False
或None
。此参数在 v1.12 中引入,有助于确保用户仅运行来自其信任的仓库的代码。如果为
False
,则会提示用户是否应信任该仓库。如果为
True
,则仓库将添加到信任列表中,无需显式确认即可加载。如果为
"check"
,则会针对缓存中的受信任仓库列表检查仓库。如果该列表中不存在,则行为将回退到trust_repo=False
选项。如果为
None
:这将发出警告,邀请用户将trust_repo
设置为False
、True
或"check"
。这仅出于向后兼容性而存在,将在 v2.0 中删除。
默认为
None
,最终将在 v2.0 中更改为"check"
。verbose (bool,可选) – 如果为
False
,则静音有关命中本地缓存的消息。请注意,无法静音有关首次下载的消息。默认为True
。
- 返回值
可用的可调用入口点
- 返回类型
示例
>>> entrypoints = torch.hub.list("pytorch/vision", force_reload=True)
- torch.hub.help(github, model, force_reload=False, skip_validation=False, trust_repo=None)[source]¶
显示入口点
model
的文档字符串。- 参数
github (str) – 格式为 <repo_owner/repo_name[:ref]> 的字符串,其中包含可选的 ref(标签或分支)。如果未指定
ref
,则假设默认分支为main
(如果存在),否则为master
。例如:'pytorch/vision:0.10'model (str) – 仓库的
hubconf.py
中定义的入口点名称的字符串force_reload (bool,可选) – 是否丢弃现有缓存并强制重新下载。默认为
False
。skip_validation (bool,可选) – 如果为
False
,则 torchhub 将检查由github
参数指定的 ref 是否正确属于 repo 所有者。这将向 GitHub API 发出请求;可以通过设置GITHUB_TOKEN
环境变量来指定非默认的 GitHub token。默认为False
。trust_repo (bool,str 或 None) –
"check"
、True
、False
或None
。此参数在 v1.12 中引入,有助于确保用户仅运行来自其信任的仓库的代码。如果为
False
,则会提示用户是否应信任该仓库。如果为
True
,则仓库将添加到信任列表中,无需显式确认即可加载。如果为
"check"
,则会针对缓存中的受信任仓库列表检查仓库。如果该列表中不存在,则行为将回退到trust_repo=False
选项。如果为
None
:这将发出警告,邀请用户将trust_repo
设置为False
、True
或"check"
。这仅出于向后兼容性而存在,将在 v2.0 中删除。
默认为
None
,最终将在 v2.0 中更改为"check"
。
示例
>>> print(torch.hub.help("pytorch/vision", "resnet18", force_reload=True))
- torch.hub.load(repo_or_dir, model, *args, source='github', trust_repo=None, force_reload=False, verbose=True, skip_validation=False, **kwargs)[source]¶
从 github 仓库或本地目录加载模型。
注意:加载模型是典型的用例,但这也可用于加载其他对象,例如标记器、损失函数等。
如果
source
为 'github',则repo_or_dir
应采用repo_owner/repo_name[:ref]
的格式,其中包含可选的 ref(标签或分支)。如果
source
为 'local',则repo_or_dir
应为本地目录的路径。- 参数
repo_or_dir (str) – 如果
source
为 'github',则这应该对应于格式为repo_owner/repo_name[:ref]
的 github 仓库,其中包含可选的 ref(标签或分支),例如 'pytorch/vision:0.10'。如果未指定ref
,则假设默认分支为main
(如果存在),否则为master
。如果source
为 'local',则它应为本地目录的路径。model (str) – 仓库/目录的
hubconf.py
中定义的可调用对象(入口点)的名称。*args (可选) – 可调用对象
model
的对应参数。source (str,可选) – 'github' 或 'local'。指定如何解释
repo_or_dir
。默认为 'github'。trust_repo (bool,str 或 None) –
"check"
、True
、False
或None
。此参数在 v1.12 中引入,有助于确保用户仅运行来自其信任的仓库的代码。如果为
False
,则会提示用户是否应信任该仓库。如果为
True
,则仓库将添加到信任列表中,无需显式确认即可加载。如果为
"check"
,则会针对缓存中的受信任仓库列表检查仓库。如果该列表中不存在,则行为将回退到trust_repo=False
选项。如果为
None
:这将发出警告,邀请用户将trust_repo
设置为False
、True
或"check"
。这仅出于向后兼容性而存在,将在 v2.0 中删除。
默认为
None
,最终将在 v2.0 中更改为"check"
。force_reload (bool,可选) – 是否无条件强制重新下载 github 仓库。如果
source = 'local'
,则无效。默认为False
。verbose (bool,可选) – 如果为
False
,则静音有关命中本地缓存的消息。请注意,无法静音有关首次下载的消息。如果source = 'local'
,则无效。默认为True
。skip_validation (bool,可选) – 如果为
False
,则 torchhub 将检查由github
参数指定的 branch 或 commit 是否正确属于 repo 所有者。这将向 GitHub API 发出请求;可以通过设置GITHUB_TOKEN
环境变量来指定非默认的 GitHub token。默认为False
。**kwargs (可选) – 可调用对象
model
的对应关键字参数。
- 返回值
使用给定的
*args
和**kwargs
调用model
可调用对象时的输出。
示例
>>> # from a github repo >>> repo = "pytorch/vision" >>> model = torch.hub.load( ... repo, "resnet50", weights="ResNet50_Weights.IMAGENET1K_V1" ... ) >>> # from a local directory >>> path = "/some/local/path/pytorch/vision" >>> model = torch.hub.load(path, "resnet50", weights="ResNet50_Weights.DEFAULT")
- torch.hub.download_url_to_file(url, dst, hash_prefix=None, progress=True)[source]¶
将给定 URL 上的对象下载到本地路径。
- 参数
示例
>>> torch.hub.download_url_to_file( ... "https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth", ... "/tmp/temporary_file", ... )
- torch.hub.load_state_dict_from_url(url, model_dir=None, map_location=None, progress=True, check_hash=False, file_name=None, weights_only=False)[source]¶
从给定的 URL 加载 Torch 序列化对象。
如果下载的文件是 zip 文件,它将自动解压缩。
如果对象已存在于 model_dir 中,则将其反序列化并返回。
model_dir
的默认值为<hub_dir>/checkpoints
,其中hub_dir
是由get_dir()
返回的目录。- 参数
url (str) – 要下载的对象的 URL
model_dir (str, 可选) – 保存对象的目录
map_location (可选) – 指定如何重新映射存储位置的函数或字典(参见 torch.load)
progress (bool, 可选) – 是否显示进度条到 stderr。默认值:True
check_hash (bool, 可选) – 如果为 True,则 URL 的文件名部分应遵循命名约定
filename-<sha256>.ext
,其中<sha256>
是文件内容的 SHA256 哈希的前八位或更多数字。哈希用于确保唯一名称并验证文件内容。默认值:Falsefile_name (str, 可选) – 下载文件的名称。如果未设置,将使用
url
中的文件名。weights_only (bool, 可选) – 如果为 True,则仅加载权重,不加载复杂的 pickle 对象。对于不可信来源建议使用。有关更多详细信息,请参见
load()
。
- 返回类型
示例
>>> state_dict = torch.hub.load_state_dict_from_url( ... "https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth" ... )
运行加载的模型:¶
请注意,torch.hub.load()
中的 *args
和 **kwargs
用于 **实例化** 模型。加载模型后,如何了解可以对模型执行哪些操作?建议的工作流程是
dir(model)
以查看模型的所有可用方法。help(model.foo)
以检查model.foo
采用哪些参数来运行
为了帮助用户在无需反复参考文档的情况下进行探索,我们强烈建议代码库所有者使函数帮助信息清晰简洁。包含一个最小的工作示例也很有帮助。
我的下载模型保存在哪里?¶
这些位置按以下顺序使用:
调用
hub.set_dir(<PATH_TO_HUB_DIR>)
$TORCH_HOME/hub
,如果设置了环境变量TORCH_HOME
。$XDG_CACHE_HOME/torch/hub
,如果设置了环境变量XDG_CACHE_HOME
。~/.cache/torch/hub
缓存逻辑¶
默认情况下,我们不会在加载文件后清理文件。如果缓存已存在于由 get_dir()
返回的目录中,则 Hub 默认使用缓存。
用户可以通过调用 hub.load(..., force_reload=True)
强制重新加载。这将删除现有的 GitHub 文件夹和下载的权重,重新初始化一个新的下载。当对同一分支发布更新时,这很有用,用户可以随时了解最新版本。
已知限制:¶
Torch hub 通过导入包来工作,就像它已安装一样。Python 中的导入会带来一些副作用。例如,您可以在 Python 缓存 sys.modules
和 sys.path_importer_cache
中看到新项,这是正常的 Python 行为。这也意味着,如果来自不同代码库的模型具有相同的子包名称(通常是 model
子包),则在从不同代码库导入不同模型时可能会出现导入错误。解决此类导入错误的一种方法是从 sys.modules
字典中删除有问题的子包;更多详细信息可以在 此 GitHub 问题 中找到。
这里值得一提的一个已知限制:用户 **不能** 在 **同一个 Python 进程** 中加载同一个代码库的两个不同分支。这就像在 Python 中安装了两个同名的包一样,这不是一个好主意。如果实际尝试这样做,缓存可能会加入进来并给您带来意外。当然,在单独的进程中加载它们是完全可以的。