Repository: pytorch/vision

Issue:
Select subset of classes to sample from
## 🚀 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
https://github.com/pytorch/vision/blob/20a771e5143c6867eee63868c38a5bcc272a35e7/torchvision/datasets/folder.py#L61

to 

```python
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


Discussion:
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 ?
@NicolasHug 
like they do here? https://discuss.pytorch.org/t/how-to-sample-images-belonging-to-particular-classes/43776/9
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
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 https://github.com/pytorch/vision/pull/3496#pullrequestreview-619539316 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 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?
> 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
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?
would it make sense to add a public method while keeping also the private one?

```python
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

```python
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?
Since the default implementation only calls the `find_classes()` function

https://github.com/pytorch/vision/blob/37eb37a836fbc2c26197dfaf76d2a3f4f39f15df/torchvision/datasets/folder.py#L205-L207

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.


> 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()`

Resolving patch:
diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py
--- a/torchvision/datasets/folder.py
+++ b/torchvision/datasets/folder.py
@@ -182,7 +182,7 @@ def __init__(
     ) -> None:
         super(DatasetFolder, self).__init__(root, transform=transform,
                                             target_transform=target_transform)
-        classes, class_to_idx = self._find_classes(self.root)
+        classes, class_to_idx = self.find_classes(self.root)
         samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)
 
         self.loader = loader
@@ -202,8 +202,12 @@ def make_dataset(
     ) -> List[Tuple[str, int]]:
         return make_dataset(directory, class_to_idx, extensions=extensions, is_valid_file=is_valid_file)
 
-    @staticmethod
-    def _find_classes(dir: str) -> Tuple[List[str], Dict[str, int]]:
+    def find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
+        """Same as :func:`find_classes`.
+
+        This method can be overridden to only consider
+        a subset of classes, or to adapt to a different dataset directory structure.
+        """
         return find_classes(dir)
 
     def __getitem__(self, index: int) -> Tuple[Any, Any]:


Test patch: