include classification labels in pretrained models

[trevoravant]

🚀 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.

Alternatives

Additional context

[asberk]

Unfortunately, the labels that I can find online don't seem to correspond to the ordering used by torchvision.models.resnet50. For instance, this example does not yield "dog" or its breed for any of the label sets I've linked to above. Granted, it gives "dog sled". But an example using my face and/or a face of a colleague produces "hook" and "bucket" as the most likely classes rather than "hair", "suit", "animal" or any number of other more suitable labels.

This is a problem because I cannot independently verify that the model is doing what it's "supposed" to be doing.

Can someone please identify a correct set of ImageNet class labels for, e.g., the ImageNet dataset that was used to train the torchvision.models.resnet50 model? I would be happy to do some of the leg work after that. (I understand that the 2012 data is no longer available from their website, but surely this is available somewhere?)

[seyeeet]

I have the same issue, did you solve it?

[asberk]

Unfortunately, no.

[t-vi]

Unfortunately, the labels that I can find online don't seem to correspond to the ordering used by torchvision.models.resnet50. For instance, this example does not yield "dog" or its breed for any of the label sets I've linked to above. Granted, it gives "dog sled". But an example using my face and/or a face of a colleague produces "hook" and "bucket" as the most likely classes rather than "hair", "suit", "animal" or any number of other more suitable labels.

I think there are two parts here:

• Having the imagenet classes in a standard way seems to be ... useful? I'm really surprised that the last issue asking for it was closed by having the classes when downloading all of imagenet (when that was still possible). It really should be readily available in TorchVision if the pretrained models give the classes.

• I have to say, the problem of predictions not matching your expectations is likely a missing model.eval(), in particular "bucket" is the typical prediction for batch size 1 and any input without .eval().

Best regards

Thomas

[datumbox]

Concerning adding the label info in the model, indeed that would be useful. We just need to consider carefully its API. Possibly this meta data would be linked to the weights and not the model architecture. We plan to tackle model versioning this half, so once this is done we would be able to safely bring this feature.

For now here is an example on how to have the labels without downloading the whole dataset: pytorch/hub#175

[t-vi]

I think we should not let the perfect generic scheme be the enemy of the pragmatic here.
How about just including the lists of category names and category ids in the ImageNet class as static members, hardcoded?
ImageNet class numbers are clearly well-standardized or the entire idea of shipping pretrained models that include the classification head would not work well.

[datumbox]

@t-vi I agree with you and it's been a problem that we haven't added it for so long. Nevertheless given that this is getting traction, it won't take too long to resolve. I was thinking that at this point we should go for a proper solution than a quick patch.

We are currently considering that labels along with other meta-data such as trained image size etc, should probably be linked with a specific version of weights as we might on the future offer pretrained weights on various datasets (say coco vs voc). I plan to pick multi-pretrained weight support (#4611) immediately after the first round of primitives of #3911 is completed.

[vadimkantorov]

I plan to pick model versioning (#2955) immediately after the first round of primitives of #3911 is completed.

Maybe until then, at least include a copy-pasteable list of well-aligned categories in some README.md file or docs online. Or maybe hard-code them in a new ImageNet dataset class? (rather than having people to use ImageFolder as in the official ImageNet training example https://github.com/pytorch/examples/blob/master/imagenet/main.py)

[datumbox]

@vadimkantorov I'm literally working on this now.

In the meanwhile, we have added a list with the ImageNet categories on hub.

You can see examples on how to easily use them for all TorchVision models here:
https://github.com/pytorch/hub/blob/4af1efa9e855389917594a62c1e4ff648ac25b61/pytorch_vision_resnet.md

[datumbox]

This information is now available on the prototype API. Hence I'll close the ticket.