facebookresearch / dino

PyTorch code for Vision Transformers training with the Self-Supervised learning method DINO
Apache License 2.0
6.06k stars 885 forks source link

normalize knn probabilities #234

Open clemsgrs opened 1 year ago

clemsgrs commented 1 year ago

In the following part of the code, the retrieved probabilities are not normalised (hence they're not really probabilities):

https://github.com/facebookresearch/dino/blob/cb711401860da580817918b9167ed73e3eef3dcf/eval_knn.py#L166-L172

I'd suggest adding a quick normalising step right after to ensure they are actually probabilities (values between 0 and 1 & summing to 1):

probs = probs / probs.sum(dim=-1).unsqueeze(-1)

it's more a "nice to have" than an actual needed change