Select subset of classes to sample from

[alemelis]

🚀 Feature

When loading a dataset with ImageFolder provide an optional argument to select a subset of classes.

Motivation

I deal with large (1000+) multi-class datasets upon which I train image classifiers. However, I usually don't want to train for all the classes at the same time.

Pitch

I'd like to change find_classes function

vision/torchvision/datasets/folder.py

Line 61 in 20a771e

classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())

to

classes = sorted(entry.name for entry in os.scandir(directory) 
                 if entry.is_dir() and (entry.name in allowed_classes or not allowed_classes))

where allowed_classes: Optional[str] = [] is an empty list by default but it can given to ImageFolder at initialisation time (it has to be propagated back to DatasetFolder where find_classes is used).

Alternatives

I tried to

1. manually create a new folder structure with only the relevant classes. This gets messy quite fast as I have now multiple versions of the same folders
2. resort to a custom dataloader filtering loaded samples after initialisation (obviously this is very slow as soon as the number of images increases)

Additional context

Here I found a discussion on the same topic.

https://discuss.pytorch.org/t/how-to-sample-images-belonging-to-particular-classes/43776/8

cc @pmeier

[NicolasHug]

Instead of adding a new allowed_classes parameter which is quite specific, would it be reasonable to instead allow overriding the self._find_classes method (by e.g. making it public, or accepting a find_classes callable parameter)?

Thoughts @pmeier ?

[alemelis]

@NicolasHug
like they do here? https://discuss.pytorch.org/t/how-to-sample-images-belonging-to-particular-classes/43776/9

[NicolasHug]

Yes, although they're overridding ImageFolder._find_classes which is private (leading underscore). Since it's private, you have no guarantee that this will work in the future: private APIs have no backward compatibility guarantees. What I'm suggesting is to "officially" support this as part of a public API - by either making _find_classes public or by allowing a new parameter

[pmeier]

I think both proposals are reasonable. We only recently factored out find_classes (#3496) but I think having a public method is a good idea. We may also address #3496 (review) in the process.

Still, the purpose of factoring out find_classes and earlier make_dataset (#3215) was to also be able to use this functionality elsewhere. I can imagine uses cases where not all classes should be loaded. Without something like allowed_classes this would be painful to do. IMO an inverted logic, i.e. excluded_classes would probably make more sense. @alemelis could you make your case for allowed_classes?

[alemelis]

@alemelis could you make your case for allowed_classes?

I'd say allowed_classes makes sense when you need to use few out of many (say 10 classes from imagenet)

on the other hand I reckon excluded_classes is handy when using most of the classes in the dataset (e.g., all the digits except "5" in MNIST)

possibly it makes sense to have both?

[NicolasHug]

possibly it makes sense to have both?

I would agree, hence my proposal about just allowing to override the callable - users can implement whatever strategy with just a little bit of code, and they can even adapt to any dataset structure

[pmeier]

After a bit more thought, IMO allowed_classes in the find_classes function makes little sense. If you already know which classes you want there is no point in calling this function. On the other hand an exclude parameter could make sense there. @alemelis since you seem to need to a known set of classes, I agree with @NicolasHug the best way to achieve it would be to overwrite ImageFolder.find_classes.

Do you want to send a PR?

[alemelis]

would it make sense to add a public method while keeping also the private one?

def _find_classes(dir: str) -> Tuple[List[str], Dict[str, int]]:
        return find_classes(dir)

def find_classes(dir: str) -> Tuple[List[str], Dict[str, int]]:
        return _find_classes(dir)

so that when overriding find_classes I could still re-use the original implementation

def new_find_classes(self, dir, allowed_classes):
        classes, class_to_idx = self.find_classes(self.root)

        # implement filtering with allowed_classes

        return classes, class_to_idx

ImageFolder.find_classes = new_find_classes

or is there a better way of doing it?

[pmeier]

Since the default implementation only calls the find_classes() function

vision/torchvision/datasets/folder.py

Lines 205 to 207 in 37eb37a

	@staticmethod
	def _find_classes(dir: str) -> Tuple[List[str], Dict[str, int]]:
	return find_classes(dir)

IMO there is no need to keep the private one. We can simply make it public. Every change to the defaults should go into the function.

[NicolasHug]

so that when overriding find_classes I could still re-use the original implementation

I agree with @pmeier, that's not really an issue right now as find_classes is public. But if the method ever becomes more than that, you could always use super()