Easily list and initialize models with new APIs in TorchVision

由瓦西利斯·弗里尼奥蒂斯和劳伦斯·鲁耶内尔编写

TorchVision 现在支持通过名称列出和初始化所有可用的内置模型和权重。这个新的 API 建立在最近引入的多权重支持 API 之上，目前处于测试阶段，并解决了社区长期以来的一个请求。

您可以在 TorchVision 的最新夜间版本中尝试这个新 API。我们正在收集反馈，以便在 TorchVision v0.14 中最终确定该功能。我们已创建一个专门的 Github Issue，您可以在那里发表评论、提问和建议！

查询和初始化可用模型

在新的模型注册 API 之前，开发者必须查询模块的 __dict__ 属性，以便列出所有可用的模型或通过名称获取特定的模型构建器方法：

# Initialize a model by its name:
model = torchvision.models.__dict__[model_name]()

# List available models:
available_models = [
    k for k, v in torchvision.models.__dict__.items()
    if callable(v) and k[0].islower() and k[0] != "_"
]

上述方法并不总是产生预期的结果，且难以发现。例如，由于 get_weight() 方法在同一个模块下公开暴露，即使它不是一个模型，也会被包含在列表中。总的来说，减少冗余（更少的导入、更短的名字等）以及能够直接通过名称初始化模型和权重（更好的对配置、TorchHub 等的支持）是社区之前提供的反馈。为了解决这个问题，我们开发了一个模型注册 API。

一种新的方法

我们在 torchvision.models 模块下添加了 4 个新方法：

from torchvision.models import get_model, get_model_weights, get_weight, list_models

样式和命名规范与 Philip Meier 提出的 Datasets V2 API 原型机制紧密一致，旨在提供类似的用户体验。出于当前仅支持 TorchVision 内置模型的考虑，模型注册方法被有意保留为私有。

列出模型

可以通过单个函数调用列出 TorchVision 中所有可用的模型：

>>> list_models()
['alexnet', 'mobilenet_v3_large', 'mobilenet_v3_small', 'quantized_mobilenet_v3_large', ...]

要列出特定子模块的可用的模型：

>>> list_models(module=torchvision.models)
['alexnet', 'mobilenet_v3_large', 'mobilenet_v3_small', ...]
>>> list_models(module=torchvision.models.quantization)
['quantized_mobilenet_v3_large', ...]

初始化模型

现在您已经知道了哪些模型可用，您可以轻松地使用预训练权重初始化一个模型：

>>> get_model("quantized_mobilenet_v3_large", weights="DEFAULT")
QuantizableMobileNetV3(
  (features): Sequential(
   ....
   )
)

获取权重

有时，在处理配置文件或使用 TorchHub 时，您可能知道特定的权重条目名称，并希望获取其实例。这可以通过以下方法轻松完成：

>>> get_weight("ResNet50_Weights.IMAGENET1K_V2")
ResNet50_Weights.IMAGENET1K_V2

要获取特定模型的全部可用权重的枚举类，您可以使用其名称：

>>> get_model_weights("quantized_mobilenet_v3_large")
<enum 'MobileNet_V3_Large_QuantizedWeights'>

或者使用其模型构建方法：

>>> get_model_weights(torchvision.models.quantization.mobilenet_v3_large)
<enum 'MobileNet_V3_Large_QuantizedWeights'>

支持 TorchHub

新方法也通过 TorchHub 提供：

import torch

# Fetching a specific weight entry by its name:
weights = torch.hub.load("pytorch/vision", "get_weight", weights="ResNet50_Weights.IMAGENET1K_V2")

# Fetching the weights enum class to list all available entries:
weight_enum = torch.hub.load("pytorch/vision", "get_model_weights", name="resnet50")
print([weight for weight in weight_enum])

把所有东西放在一起

例如，如果您想检索所有具有预训练权重的中小型模型并初始化其中一个，只需使用上述 API 即可：

import torchvision
from torchvision.models import get_model, get_model_weights, list_models


max_params = 5000000

tiny_models = []
for model_name in list_models(module=torchvision.models):
    weights_enum = get_model_weights(model_name)
    if len([w for w in weights_enum if w.meta["num_params"] <= max_params]) > 0:
        tiny_models.append(model_name)

print(tiny_models)
# ['mnasnet0_5', 'mnasnet0_75', 'mnasnet1_0', 'mobilenet_v2', ...]

model = get_model(tiny_models[0], weights="DEFAULT")
print(sum(x.numel() for x in model.state_dict().values()))
# 2239188

更多技术细节请参阅原始 RFC。请抽出几分钟时间对新的 API 提供反馈，这对将其从测试版毕业并包含在下一个版本中至关重要。您可以在专门的 Github Issue 上这样做。我们期待阅读您的评论！

使用新的 API 轻松列出和初始化模型