Classification with trained model: Dataset __gettitem__ override?

I have a trained model that I want to use it to classify images. These images, since unclassified, obviously do not follow the format required by Imageloader/Dataset. One hack is to simulate the nested structure required and move all images to a single subfolder, but this doesn’t work for large datasets. Others suggest a custom dataset but I’m unclear how exact to override the Dataset class with my own __getitem__ method to achieve this.

Since this is such a standard thing (in production, for testing, etc.) I’m surprised there is no straightforward way, or I’m likely just not finding it.

Any pointers appreciated. Thanks.

This tutorial should give you a simple example how to write a custom Dataset.
The __getitem__ method will accept an index in the standard use case, which is used to load, process, and return the sample (and target if possible).
You can use other libraries such as PIL to load the image and apply any preprocessing on this sample you like.
The DataLoader will then create a batch of all returned samples for the current indices.

1 Like

Thanks, just out of curiosity, I’m curious to hear if you have any opinion about why this is not an out-of-the-box feature? Granted, it’s not difficult to implement but isn’t this a basic feature for a framework like this? Could it be the more academic “bias” (as opposed to industry like TF) of pytorch? Cheers.

I’m not sure which feature you would like to see.
Would you like to use a ImageFolder-like dataset, where you could specify the targets manually (or reassign them somehow)?
Or do you think the custom Dataset implementation is too complicated?

If I have a trained model and want to use it in production to classify new, unlabeled images in some directory, right now I cannot use an out of the box utility to load them for transformation.

As the API works right now, a user needs to either a) move all the images to a new subdirectory to conform to the nested dir structure expected by ImageFolder or b) write some custom function as in the tutorial you provide.

If you see my github example above you will see what I mean. The issue was opened in October 2017, and the last response of a user with similar needs was April 15, 2020. Clearly a subset of users need this basic functionality. Moving millions of images every day from the root dir to a subdir to have the nested structure is a waste of time and server moneys.

Having every user do either a) or b) above seems silly imo for something that is arguably an elementary purpose of having a trained model, which is actually using it to classify unlabeled images as part of a bigger pipeline.

Something like this would be useful:

imgs = datasets.FlatImageFolder(

where non_nested_image_dir is simply a dir with the images to be classified inside. Bear in mind that these images have no labels --that’s the purpose of the trained network.

Again, its not that writing a custom function is hard, but the fact that it seems such an ovious thing to have in the API and some users clearly need it (I suspect practitioners building pipelines). Hope this helps clarify what I meant.


That’s right and I would recommend to use symbolic links instead of moving data around, if necessary.

That makes sense and we would be happy to accept a PR for it. :slight_smile:
I would suggest to either post in the linked issue or create a new feature request in the torchvision repository to discuss the implementation.

Will do, thanks for the pointer.

1 Like