从 torchvision 预训练模型中获取模型类标签

Question

我正在使用来自 torchvision 的预训练 Alexnet 模型（没有微调）。问题是，即使我能够在某些数据上运行模型并获得输出概率分布，我也无法找到类标签将其映射到.

按照这个官方文档

import torch

model = torch.hub.load('pytorch/vision:v0.6.0', 'alexnet', pretrained=True)

model.eval()

AlexNet(

  (features): Sequential(

    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))

    (1): ReLU(inplace=True)

    (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)

    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))

    (4): ReLU(inplace=True)

    (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)

    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

    (7): ReLU(inplace=True)

    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

    (9): ReLU(inplace=True)

    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

    (11): ReLU(inplace=True)

    (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)

  )

  (avgpool): AdaptiveAvgPool2d(output_size=(6, 6))

  (classifier): Sequential(

    (0): Dropout(p=0.5, inplace=False)

    (1): Linear(in_features=9216, out_features=4096, bias=True)

    (2): ReLU(inplace=True)

    (3): Dropout(p=0.5, inplace=False)

    (4): Linear(in_features=4096, out_features=4096, bias=True)

    (5): ReLU(inplace=True)

    (6): Linear(in_features=4096, out_features=1000, bias=True)

  )

)

在处理图像的一些步骤之后，我能够使用它来获得单个图像的输出作为 (1,1000) dim 向量，我将使用 softmax 来获得概率分布 -

Answer 1

不幸的是，您不能直接从 torchvision 模型中获取类标签名称。但是，这些模型是在 ImageNet 数据集（因此有 1000 个类）上训练的。据我所知，您必须从网上获取类名映射；没有办法把它从火炬上拿下来。以前，您可以使用 torchvision.datasets.ImageNet 直接下载 ImageNet，它有一个内置的标签到类名转换器。现在下载链接不公开，需要手动下载，然后才能被 datasets.ImageNet 使用。因此，您可以简单地在线搜索 ImageNet 的类到标签映射，而不是下载数据或尝试使用 torch。