通过 PyTorch Hub 实现可复现的研究 – PyTorch

可复现性是许多研究领域（包括基于机器学习技术的研究）的基本要求。然而，许多机器学习领域的出版物要么无法复现，要么难以复现。随着研究出版物数量的持续增长（包括 arXiv 上托管的数万篇论文以及创历史新高的会议投稿量），研究的可复现性比以往任何时候都更加重要。尽管许多出版物都附带了代码和预训练模型，这很有帮助，但用户仍然需要自行摸索许多步骤才能实现复现。

我们很高兴地宣布推出 PyTorch Hub，这是一个简单的 API 和工作流程，为提高机器学习研究的可复现性提供了基础构件。PyTorch Hub 由一个专门为促进研究可复现性并赋能新研究而设计的预训练模型库组成。它还内置了对 Colab 的支持，与 Papers With Code 进行了集成，目前包含了广泛的模型集，包括分类与分割、生成模型、Transformers 等。

[所有者] 发布模型

PyTorch Hub 支持通过添加一个简单的 hubconf.py 文件，将预训练模型（模型定义和预训练权重）发布到 GitHub 仓库。这提供了对所支持模型的枚举以及运行模型所需依赖项的列表。示例可以在 torchvision、huggingface-bert 和 gan-model-zoo 仓库中找到。

让我们看看最简单的情况：torchvision 的 hubconf.py

# Optional list of dependencies required by the package
dependencies = ['torch']

from torchvision.models.alexnet import alexnet
from torchvision.models.densenet import densenet121, densenet169, densenet201, densenet161
from torchvision.models.inception import inception_v3
from torchvision.models.resnet import resnet18, resnet34, resnet50, resnet101, resnet152,\
resnext50_32x4d, resnext101_32x8d
from torchvision.models.squeezenet import squeezenet1_0, squeezenet1_1
from torchvision.models.vgg import vgg11, vgg13, vgg16, vgg19, vgg11_bn, vgg13_bn, vgg16_bn, vgg19_bn
from torchvision.models.segmentation import fcn_resnet101, deeplabv3_resnet101
from torchvision.models.googlenet import googlenet
from torchvision.models.shufflenetv2 import shufflenet_v2_x0_5, shufflenet_v2_x1_0
from torchvision.models.mobilenet import mobilenet_v2

在 torchvision 中，模型具有以下属性：

每个模型文件都可以独立运行和执行
它们除了 PyTorch 之外不需要任何其他软件包（在 hubconf.py 中编码为 dependencies['torch']）
它们不需要单独的入口点，因为模型在创建时就可以直接开箱即用

最小化软件包依赖可以降低用户加载您的模型并进行即时实验的阻碍。

一个更复杂的例子是 HuggingFace 的 BERT 模型。这是他们的 hubconf.py

dependencies = ['torch', 'tqdm', 'boto3', 'requests', 'regex']

from hubconfs.bert_hubconf import (
    bertTokenizer,
    bertModel,
    bertForNextSentencePrediction,
    bertForPreTraining,
    bertForMaskedLM,
    bertForSequenceClassification,
    bertForMultipleChoice,
    bertForQuestionAnswering,
    bertForTokenClassification
)

每个模型都需要创建一个入口点。以下是指定 bertForMaskedLM 模型入口点的代码片段，该模型返回预训练的模型权重。

def bertForMaskedLM(*args, **kwargs):
    """
    BertForMaskedLM includes the BertModel Transformer followed by the
    pre-trained masked language modeling head.
    Example:
      ...
    """
    model = BertForMaskedLM.from_pretrained(*args, **kwargs)
    return model

这些入口点可以作为复杂模型工厂的包装器。它们可以提供清晰且一致的帮助文档字符串，包含支持下载预训练权重的逻辑（例如通过 pretrained=True），或者具有额外的 Hub 特定功能，如可视化。

有了 hubconf.py，您就可以根据此处的模板提交拉取请求 (Pull Request)。我们的目标是策划高质量、易于复现且对研究可复现性最有益的模型。因此，我们可能会与您合作完善您的拉取请求，在某些情况下可能会拒绝发布一些低质量的模型。一旦我们接受您的拉取请求，您的模型很快就会出现在 PyTorch Hub 网页上，供所有用户探索。

[用户] 工作流程

作为用户，PyTorch Hub 让您只需遵循几个简单的步骤即可完成以下操作：1) 探索可用模型；2) 加载模型；3) 了解给定模型有哪些可用方法。让我们分别浏览一下这些示例。

探索可用入口点。

用户可以使用 torch.hub.list() API 列出仓库中所有可用的入口点。

>>> torch.hub.list('pytorch/vision')
>>>
['alexnet',
'deeplabv3_resnet101',
'densenet121',
...
'vgg16',
'vgg16_bn',
'vgg19',
 'vgg19_bn']

请注意，PyTorch Hub 也允许辅助入口点（预训练模型之外的），例如用于 BERT 模型预处理的 bertTokenizer，以使用户的工作流程更加顺畅。

加载模型

现在我们知道了 Hub 中有哪些模型可用，用户可以使用 torch.hub.load() API 加载模型入口点。这只需要一条命令，无需安装 wheel 包。此外，torch.hub.help() API 可以提供有关如何实例化模型的有用信息。

print(torch.hub.help('pytorch/vision', 'deeplabv3_resnet101'))
model = torch.hub.load('pytorch/vision', 'deeplabv3_resnet101', pretrained=True)

仓库所有者通常也希望持续添加错误修复或性能改进。PyTorch Hub 让用户可以通过调用以下命令非常轻松地获取最新更新：

model = torch.hub.load(..., force_reload=True)

我们相信这将有助于减轻仓库所有者重复发布软件包的负担，从而让他们能够更专注于研究。这也确保了作为用户，您可以获得最新可用的模型。

相反，稳定性对用户来说很重要。因此，一些模型所有者会从特定的分支或标签（而不是 master 分支）提供模型，以确保代码的稳定性。例如，pytorch_GAN_zoo 从 hub 分支提供模型。

model = torch.hub.load('facebookresearch/pytorch_GAN_zoo:hub', 'DCGAN', pretrained=True, useGPU=False)

请注意，传递给 hub.load() 的 *args 和 **kwargs 用于实例化模型。在上面的示例中，pretrained=True 和 useGPU=False 被传递给了模型的入口点。

探索已加载的模型

一旦您从 PyTorch Hub 加载了模型，您可以使用以下工作流程来找出支持的可用方法，并更好地了解运行它需要哪些参数。

使用 dir(model) 查看模型的所有可用方法。让我们看看 bertForMaskedLM 的可用方法。

>>> dir(model)
>>>
['forward'
...
'to'
'state_dict',
]

help(model.forward) 提供了关于运行已加载模型所需参数的视图。

>>> help(model.forward)
>>>
Help on method forward in module pytorch_pretrained_bert.modeling:
forward(input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None)
...

请仔细查看 BERT 和 DeepLabV3 页面，在那里您可以了解模型加载后如何使用。

其他探索方式

PyTorch Hub 中可用的模型也支持 Colab，并且直接链接到 Papers With Code，您可以一键开始使用。此处是一个很好的入门示例（如下所示）。

其他资源：

PyTorch Hub API 文档可以在此处找到。
在此提交模型以在 PyTorch Hub 上发布。
前往 https://pytorch.ac.cn/hub 了解更多关于可用模型的信息。
在 paperswithcode.com 上查看即将推出的更多模型。

非常感谢 HuggingFace 的同仁、PapersWithCode 团队、fast.ai 和 Nvidia，以及 Morgane Riviere (FAIR Paris) 和其他许多人为启动这项工作所提供的帮助！！

干杯！

PyTorch 团队

常见问题解答：

问：如果我们想贡献一个 Hub 中已有的模型，但我的模型准确率更高，我还应该贡献吗？

答：应该！！Hub 的下一步是实现投票系统以筛选出最好的模型。

问：谁来托管 PyTorch Hub 的模型权重？

答：作为贡献者，您有责任托管模型权重。您可以将模型托管在您喜欢的云存储中，或者如果它符合限制，可以托管在 GitHub 上。如果您无法负担托管权重的成本，请通过在 hub 仓库中开启 issue 与我们联系。

问：如果我的模型是在私有数据上训练的，我还能贡献这个模型吗？

答：不能！PyTorch Hub 以开放研究为核心，这延伸到使用开放数据集来训练这些模型。如果提交了专有模型的拉取请求，我们将友好地请您重新提交一个在开放且可用数据集上训练的模型。

问：我下载的模型保存在哪里？

答：我们遵循 XDG 基础目录规范，并遵守有关缓存文件和目录的通用标准。

位置按以下顺序使用：

调用 hub.set_dir(<PATH_TO_HUB_DIR>)
$TORCH_HOME/hub（如果设置了环境变量 TORCH_HOME）
$XDG_CACHE_HOME/torch/hub（如果设置了环境变量 XDG_CACHE_HOME）
~/.cache/torch/hub

使用 PyTorch Hub 实现可复现研究