vturrisi / solo-learn

solo-learn: a library of self-supervised methods for visual representation learning powered by Pytorch Lightning
MIT License
1.41k stars 184 forks source link

Training on custom dataset #237

Closed AlexanderUzhinskiy closed 2 years ago

AlexanderUzhinskiy commented 2 years ago

Hello! First of all - thank you for your work! You have a very interesting project!

I'm trying to use it to train a base network on custom data for our project related to plant diseases detection.

I managed to do it using byol approach. In the end, I have byol-400ep-custom-3u5w2uqi-ep=40.ckpt file. If I understand right - I can use it as a feature extractor.

There is no known number of classes just a huge amount of plant images. So, I want to know is there a way to understand how good the network is in classification. Can I somehow get suggestions on the number of classes and see examples of images related to classes?

Is there a way to use mobilenet_v2 as the backbone?

vturrisi commented 2 years ago

Hey, Thanks. You can try clustering with k-means to discover some properties of your data. Or manually label some instances and perform k-NN (then you can use something like our main_knn.py). There's no property of BYOL or our library that would suggest a number of possible clusters in the data.

About the backbone, sure, but you will need to retrain. You will also need to modify the library, adding the backbone that you want to solo/utils/backbones.py and add it as an option here: https://github.com/vturrisi/solo-learn/blob/d5ac12d3954184f79d93dee6f52b67da1408afd2/solo/methods/base.py#L69-L89

Note that backbone-specific parameters need to be passed to the model, modifying the code here: https://github.com/vturrisi/solo-learn/blob/d5ac12d3954184f79d93dee6f52b67da1408afd2/solo/methods/base.py#L285-L292 and in solo/args/utils.py and also parsing everything in the init of the base method class.

AlexanderUzhinskiy commented 2 years ago

Thank you! I will try!