получение классификационных ярлыков для предварительно обученных сетей Torchvision

Question

Пакет Pytorch torchvision предоставляет предварительно обученные нейронные сети для классификации изображений. Я использовал следующий код для классификации изображения с использованием Ale xnet (примечание: часть этого кода взята с этой веб-страницы ):

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)

Всего 1000 классов и переменная top5_inds дает мне индексы из 5 лучших классов. Но как я могу получить соответствующие ярлыки (например, улитка, баскетбол, банан)? Я не могу найти какой-либо список как часть документации Pytorch или переменной alexnet.

Answer 1

Модели Torchvision проходят предварительную подготовку по набору данных Imag eNet. Благодаря своей полноте и размеру Imag eNet является наиболее часто используемым набором данных для подготовки и обучения. Как вы заметили, у него 1000 классов. Полный список классов можно найти, или вы можете обратиться к этому списку на GitHub: https://gist.github.com/yrevar/942d3a0ac09ec9e5eb3a