Decouple feature extractor from models and wrap the dataset instead

[jpcbertoldo]

Context

Several methods have the feature extraction as the first step of the pipeline, and (unless your dataset uses data augmentation) a feature vector will always be the same given an image and a pre-trained model.

Idea

I propose to make a wrapper for (not sure which or all) DataLoader/DataSet/DataModule (DL/DS/DM), most likely, with a decorator
that carries a (frozen? for the sake of reproducibility and simplicity) feature extractor network inside.

Its role is to extract the feature and pass it along so the models don't have to implement this.

Downstream features

With such mechanism, both features below would be possible/available to any model.

1) cache the feature vectors [by image idx] to speed up training

A simple cache mechanism like functools.cache would make the feature extraction faster during trainings that iterate over the dataset multiple times.

• Complication: augmentations may make this challenging.
• Solution 1 (easy): feature only available when there are no augmentations.
• Solution 2 (hard): keep track of augmentation params and design cache keys using the combination of transformation params with the image index/path -- feasible if using only discrete random transform parameters.

2) Augmentations in the feature space

One can play both with augmentations in the image space (as it's done with the transform_config now) and then on the feature space as well. And this could also be eventually cached.

[ashwinvaidya17]

I like this idea. We won't have to pass the image through the feature extractor each time the same image is passed to the model during training. I am concerned about the memory usage but we can think of some design that mitigates that.

[jpcbertoldo]

I am concerned about the memory usage

I'd just leave it up to the user to deal with that through an option.

The decorator functools.cache has an argument maxsize, which is the number of calls cached in the memory.

If something fancier based rather on memory megabytes is necessary it shouldn't be all that complicated to implement i guess.

[ashwinvaidya17]

A few things to keep in mind would be:

1. How the model export and inference scripts are affected by it. Currently, the feature extractor is "baked" into the models when we export to onnx and openvino. If we do it from dataloader then the model will expect n feature maps as the input. I am not sure if we will be able to call model(image) in our inference scripts. Might affect deployment. But caching is something we can definitely add to the training step within the torch model and check with self.training.
2. How is tiling affected by it. Currently the tiler object is added to the torch model so that it can be exported as well. If we do it from within the dataloader then we might have to move the tiler as well. Not sure if it will be the cleanest of refactor.

How is this alternative. We can add the decorator to torch model's forward method. This way we can pass an image to it and feature extractor will call tiler and return the features. The LRU cache can use image hash as keys and avoid the whole computation. This way will be able to preserve feature extraction and tiling in the model graph.

[jpcbertoldo]

I have no idea of how export and tilling works, I'll trust on you on this one : )

We can add the decorator to torch model's forward method.

Sounds right as well.

[samet-akcay]

How is this alternative. We can add the decorator to torch model's forward method. This way we can pass an image to it and feature extractor will call tiler and return the features. The LRU cache can use image hash as keys and avoid the whole computation. This way will be able to preserve feature extraction and tiling in the model graph.

I guess this could be taken into account when working on the new tiling design

[samet-akcay]

One more item to consider is how this is done for custom datasets? It would probably be straightforward for MVTec and other benchmark datasets, but not sure how this would be handled for custom datasets.

[jpcbertoldo]

One more item to consider is how this is done for custom datasets? It would probably be straightforward for MVTec and other benchmark datasets, but not sure how this would be handled for custom datasets.

What kind of problem are you thinking about?

As it would be a decorator, it should be transparent to whoever is using it and the decorated object should be unaware.
So, by design it is supposed to be adapted to any dataset a priori.