image-processing - 获取 torchvision 预训练网络的分类标签

Question

Pytorch 的 torchvision 包提供了用于图像分类的 pre-trained neural networks。我一直在使用以下代码使用 Alexnet 对图像进行分类(注意:部分代码来自 this webpage ):

from PIL import Image
import torch
from torchvision import transforms
from torchvision import models

# function to transform image
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225])])

# image
img = Image.open('/path/to/image.jpg')
img = transform(img)
img = torch.unsqueeze(img, 0)

# alexnet
alexnet = models.alexnet(pretrained=True)
alexnet.eval()
out = alexnet(img)
percents = torch.nn.functional.softmax(out, dim=1)[0] * 100
top5_vals, top5_inds = percents.topk(5)

总共有 1,000 个类，top5_inds 变量为我提供了前 5 个类的索引。但是我如何获得相关标签(例如蜗牛、篮球、香蕉)？我似乎找不到任何类型的列表作为 Pytorch 文档或 alexnet 变量的一部分。

Answer 1

Torchvision 模型在 ImageNet 数据集上进行了预训练。由于其全面性和规模，ImageNet 是最常用于预训练和迁移学习的数据集。如您所述，它有 1000 个类。完整的类列表可以搜索，也可以引用GitHub上的这个列表:https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a