使用 PyTorch Hub 实现可复现研究

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

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

[所有者] 发布模型

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，您可以根据 此处 的模板发送拉取请求。我们的目标是为研究可复现性整理高质量、易于复现、最大化受益的模型。因此，我们可能会与您合作完善您的拉取请求，在某些情况下可能会拒绝发布一些低质量模型。一旦我们接受您的拉取请求，您的模型将很快出现在 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 仓库上提出问题与我们联系。

问：如果我的模型是在私有数据上训练的怎么办？我还应该贡献这个模型吗？

答：不！PyTorch Hub 围绕开放研究，这延伸到使用开放数据集来训练这些模型。如果提交了专有模型的拉取请求，我们将恳请您重新提交一个在开放和可用数据上训练的模型。

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

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

位置按以下顺序使用

• 调用 hub.set_dir(<PATH_TO_HUB_DIR>)
• 如果设置了环境变量 TORCH_HOME，则为 $TORCH_HOME/hub。
• 如果设置了环境变量 XDG_CACHE_HOME，则为 $XDG_CACHE_HOME/torch/hub。
• ~/.cache/torch/hub