Should batchnorm layers be updated after avg_checkpoints.py?

[hal-314]

Hello,

Reading the new Pytorch post about SWA released in 1.6, they recommend to update batch norm statistic after averaging the checkpoints. I notice that avg_checkpoints.py doesn't do it. With the reason that Pytorch team for updating bn statistics, it makes sense that avg_checkpoints.py does it.

The new torch.optim.swa_utils.update_bn is pretty straigh forward. So, it could be copied (I don't know if there is a license conflict) to avoid forcing version 1.6. If you are considering to upgrade to Pytorch 1.6, we could use the new AverageModel class to simplify the script.

Finally, I'm willing to make a PR for both approaches.

[rwightman]

@hal-314 I've tried this before, and it made things worse. So to keep it simple I just leave BN layers as part of the EMA averaging and also the avg_checkpoints script. Unlike SWA, the weights avgeraged with avg_checkpoints likely don't differ quite as much as I pick off the N best, which tend to cluster towards the end of my usual training schedule I don't crank up the LR at the end like SWA.

I'm all for someone proving me wrong, I just don't consider this my highest priority. So please, anyone is welcome to try averaging checkpoints with and without updating the BN stats and report results here.

[pgsrv]

Hi! Is there an example code of updating bn stats when running avg_checkpoints?

[JeanKaddour]

I observed that updating BN stats in a ResNet improved the performance a tiny bit, like ~0.25%.
More details can be found here: https://arxiv.org/pdf/2209.14981.pdf