Using ViTForClassification for regression?


I would like to use the ViT model for classification and adapt it to a regression task, is it feasible ?

Can the model work just by changing the loss function ? How can I define the classes in the _info method of my custom dataset since there is an infinity of them possible ? What are all the other changes to make ?

Thank you

If you set the num_labels of the config to 1, it will automatically use the MSE loss for regression, as can be seen here. So yes, it’s totally possible.