search

augustwester / vicreg

A PyTorch implementation of VICReg by Bardes et al. (2021), trained on CIFAR-10 with a ResNet-18 backbone.
https://sigmoidprime.com/post/vicreg/
MIT License
6 stars 0 forks source link

needs a model.eval() #1

Closed neel-dey closed 1 year ago

neel-dey commented 1 year ago

Thanks for the minimal VICReg implementation!

eval.py is missing a model.eval() or model.encoder.eval() call prior to the linear layer training loop. Without it, the BatchNorm moving statistics are still being updated in the encoder network with each forward pass and it's not doing linear probing.

As expected, using eval mode and linear probing only brings performance down from 86.99% to 85.50%.

Best, Neel

augustwester commented 1 year ago

Hi Neel,

Good catch! Thanks for letting me know. Will update the code and readme later today.

augustwester commented 1 year ago

Fixed in 89eab04.