include classification labels in pretrained models 路 Issue #1946 路 pytorch/vision (original) (raw)
馃殌 Feature
Include classification labels in pretrained models
Motivation
Let's say I want to use one of Torchvision's pre-trained classification networks, e.g. torchvision.models.alexnet(pretrained=True). The AlexNet model has 1,000 outputs which correspond to specific classes (i.e. the ImageNet 1,000 classes). To the best of my knowledge, the class labels (e.g. snail, basketball, banana) are not accessible anywhere in Pytorch or Torchvision.
This means that if someone wants to use a pre-trained model, they have to go through the extra steps of finding a file containing a list of the ImageNet 1,000 classes, and then importing and parsing that file to get those class names into their code.
So I think it would be great if the class labels were natively included somewhere in Torchvision.
Pitch
My best idea is to represent the class labels as a list, which would be an attribute of the model:
from torchvision import models alexnet = models.alexnet(pretrained=True) classes = alexnet.classes classes[0] 'tench, Tinca tinca' classes[1] 'goldfish, Carassius auratus'
It may make sense to only define this list when the model is obtained using pretrained=True.