EMA Update of bn buffer

The following function apply moving average to the ema model. But it didn’t update the statistic(runing_mean and runing_var) since these two were not parameters but buffers.

def update_moving_average(ema_updater, ma_model, current_model):
    for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
        old_weight, up_weight = ma_params.data, current_params.data
        ma_params.data = ema_updater.update_average(old_weight, up_weight)

Should I use this function instead?

def update_moving_average(ema_updater, ma_model, current_model):
    for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
        old_weight, up_weight = ma_params.data, current_params.data
        ma_params.data = ema_updater.update_average(old_weight, up_weight)
    for current_buffers, ma_buffers in zip(current_model.buffers(), ma_model.buffers()):
        old_weight, up_weight = ma_buffers.data, current_buffers.data
        ma_params.data = ema_updater.update_average(old_weight, up_weight)

[lucidrains]

@luyvlei so i think the issue is because the batchnorm statistics are already a moving average - i’ll have to read the momentum squared paper above in detail and see if the conclusions are sound

as an aside, there are papers that are starting to use SimSiam (kaimings work where the teacher is the same as the student, but with a stop gradient) successfully, and which does not require exponential moving averages as does BYOL. so i’m wondering how important these little details are, and whether it is worth the time to even debug

https://arxiv.org/abs/2111.00210 https://arxiv.org/abs/2110.05208

[luyvlei]

Yep, if the target model uses train mode, the statics of BN doesn’t matter since it will never be used. And in this implementation it also use train mode. But it’s not clear that the EVAL mode would have yielded any better results @iseessel