torch.hub¶
Pytorch Hub 是一个预训练模型仓库，旨在促进研究的可复现性。
发布模型¶
Pytorch Hub 支持通过添加一个简单的 hubconf.py 文件，将预训练模型（模型定义和预训练权重）发布到 GitHub 仓库中；
hubconf.py 可以有多个入口点（entrypoint）。每个入口点定义为一个 python 函数（例如：你想发布的预训练模型）。
def entrypoint_name(*args, **kwargs):
# args & kwargs are optional, for models which take positional/keyword arguments.
...
如何实现入口点（entrypoint）？¶
以下是一个代码片段，它指定了 resnet18 模型的入口点（entrypoint），如果我们展开 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会传递给实际的可调用函数。函数的 Docstring 用作帮助消息。它解释了模型的功能以及允许的位置/关键字参数。强烈建议在此处添加一些示例。
入口点函数可以返回一个模型（nn.module），也可以返回辅助工具，以使用户工作流程更顺畅，例如分词器（tokenizers）。
以 (underscore) 前缀开头的可调用函数被视为辅助函数，它们不会在
torch.hub.list()中显示。预训练权重可以存储在 GitHub 仓库本地，也可以通过
torch.hub.load_state_dict_from_url()加载。如果小于 2GB，建议将其附加到 项目发布 (project release) 中并使用发布中的 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() 显示 Docstring 和示例，并使用 torch.hub.load() 加载预训练模型。
- torch.hub.list(github, force_reload=False, skip_validation=False, trust_repo=None, verbose=True)[source][source]¶
列出
github参数指定的仓库中所有可用的可调用入口点（callable entrypoints）。- 参数
github (str) – 格式为 “repo_owner/repo_name[:ref]” 的字符串，其中 ref (tag 或 branch) 是可选的。如果未指定
ref，则假定默认分支为main（如果存在），否则为master。示例：‘pytorch/vision:0.10’force_reload (bool, optional) – 是否丢弃现有缓存并强制重新下载。默认为
False。skip_validation (bool, optional) – 如果为
False，torchhub 将检查github参数指定的分支或提交是否正确地属于仓库所有者。这将向 GitHub API 发出请求；你可以通过设置GITHUB_TOKEN环境变量来指定非默认的 GitHub 令牌。默认为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, optional) – 如果为
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][source]¶
显示入口点
model的 Docstring。- 参数
github (str) – 格式为 <repo_owner/repo_name[:ref]> 的字符串，其中 ref (tag 或 branch) 是可选的。如果未指定
ref，则假定默认分支为main（如果存在），否则为master。示例：‘pytorch/vision:0.10’model (str) – 在仓库的
hubconf.py中定义的入口点名称的字符串。force_reload (bool, optional) – 是否丢弃现有缓存并强制重新下载。默认为
False。skip_validation (bool, optional) – 如果为
False，torchhub 将检查github参数指定的 ref 是否正确地属于仓库所有者。这将向 GitHub API 发出请求；你可以通过设置GITHUB_TOKEN环境变量来指定非默认的 GitHub 令牌。默认为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][source]¶
从 github 仓库或本地目录加载模型。
注意：加载模型是典型的用例，但这也可用于加载其他对象，如分词器（tokenizers）、损失函数等。
如果
source是 ‘github’，则repo_or_dir预期为repo_owner/repo_name[:ref]格式，其中 ref (tag 或 branch) 是可选的。如果
source是 ‘local’，则repo_or_dir预期为本地目录的路径。- 参数
repo_or_dir (str) – 如果
source是 ‘github’，这应该对应一个 github 仓库，格式为repo_owner/repo_name[:ref]，其中 ref (tag 或 branch) 是可选的，例如 ‘pytorch/vision:0.10’。如果未指定ref，则假定默认分支为main（如果存在），否则为master。如果source是 ‘local’，则它应该是一个本地目录的路径。model (str) – 在仓库/目录的
hubconf.py中定义的可调用函数（入口点）的名称。*args (optional) – 可调用函数
model对应的 args。source (str, optional) – ‘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, optional) – 是否强制无条件重新下载 github 仓库。如果
source = 'local'，此参数无效。默认为False。verbose (bool, optional) – 如果为
False，则静默有关命中本地缓存的消息。请注意，首次下载的消息无法静默。如果source = 'local'，此参数无效。默认为True。skip_validation (bool, optional) – 如果为
False，torchhub 将检查github参数指定的分支或提交是否正确地属于仓库所有者。这将向 GitHub API 发出请求；你可以通过设置GITHUB_TOKEN环境变量来指定非默认的 GitHub 令牌。默认为False。**kwargs (optional) – 可调用函数
model对应的 kwargs。
- 返回值
当使用给定的
*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][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][source]¶
加载给定 URL 处的 Torch 序列化对象。
如果下载的文件是 zip 文件，它将自动解压。
如果对象已存在于 model_dir 中，则将其反序列化并返回。
model_dir的默认值为<hub_dir>/checkpoints，其中hub_dir是由get_dir()返回的目录。- 参数
url (str) – 要下载的对象的 URL。
model_dir (str, optional) – 保存对象的目录。
map_location (optional) – 一个函数或字典，指定如何重新映射存储位置（参见 torch.load）。
progress (bool, optional) – 是否向 stderr 显示进度条。默认值：True。
check_hash (bool, optional) – 如果为 True，URL 的文件名部分应遵循命名约定
filename-<sha256>.ext，其中<sha256>是文件内容的 SHA256 哈希的前八位或更多位数字。哈希用于确保名称的唯一性并验证文件内容。默认值：False。file_name (str, optional) – 下载文件的名称。如果未设置，将使用
url中的文件名。weights_only (bool, optional) – 如果为 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，则使用$TORCH_HOME/hub。如果设置了环境变量
XDG_CACHE_HOME，则使用$XDG_CACHE_HOME/torch/hub。~/.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 issue。
一个值得在此提及的已知限制是：用户不能在同一个 Python 进程中加载同一个仓库的两个不同分支。这就像在 Python 中安装两个同名的包一样，不是一个好的做法。如果您真的尝试这样做，缓存可能会加入进来并给您带来惊喜。当然，在单独的进程中加载它们是完全没问题的。
