It would be great to train and evaluate a linear probe on the validation / test set during pretraining. That is especially useful during hyperparameter tuning / development / debugging to have an idea of how well the model is training and possibly whether to kill the job. This would also allow having linear probing evaluation without requiring a second stage (which for ImageNet is actually quite slow as the images have to be encoded at every epoch). This makes a huge difference when you don't have so much compute.
If one uses the same optimizer for the encoder and linear probe then this should be a very easy addition. Using a different optimizer might require a little more work.
🚀 Feature
It would be great to train and evaluate a linear probe on the validation / test set during pretraining. That is especially useful during hyperparameter tuning / development / debugging to have an idea of how well the model is training and possibly whether to kill the job. This would also allow having linear probing evaluation without requiring a second stage (which for ImageNet is actually quite slow as the images have to be encoded at every epoch). This makes a huge difference when you don't have so much compute.
If one uses the same optimizer for the encoder and linear probe then this should be a very easy addition. Using a different optimizer might require a little more work.
Pytorch lightning bolts have a callback for that: https://lightning-bolts.readthedocs.io/en/latest/callbacks/self_supervised.html#sslonlineevaluator