To allow batchnorm layers with frozen running stats

[EmmyFang]

🚀 Feature

Is it possible for the privacy engine to allow for batchnorm layers when we freeze their running stats (i.e., all batchnorm layers in .eval() mode)?

Motivation

When we want to use transfer learning and the pretrained model has batchnorm layers, it would be helpful if we can still use the privacy engine by freezing their running stats and treating the stats as constants.

Pitch

It would be helpful if the privacy engine allows for modules like batchnorm with frozen running stats, especially when we want to use some pretrained models with batchnorm layers. Thank you in advance!

[pierrestock]

Hey EmmyFang,

Thanks for your interest. I think that's a good idea of enhancement to include in a future release.

However, one question remains: how do we exactly initialize the (frozen) running statistics? If we use the defaults (mean to 0, std to 1), then we might as well replace this layer by a channel-wise affine transform (and either define its grad sample as usual or leverage functorch.

Pierre

[EmmyFang]

Hi Pierre,

I'm not sure if I understood your question correctly. I think using channel-wise affine transform to replace BN layers is a good idea since the key is just to make sure we can (1) load the pretrained information, and (2) the running stats will not get further updated

I'm imaging something as follows:

• initialize a model (with the BN modules) and load the pretrained weights
• for each BN module in the original model, initialize channel-wise affine transformation module, load the mean and variance as well as the weight and bias from the BN module
• replace the BN layers in the original model with the channel-wise affine transformation

The user can perform the first step to load the pretrained weights and then maybe ModuleValidator.fix() can provide the option to perform the second and third steps.

Thanks,
Emmy

[pierrestock]

Hey Emmy,

Apologies for the late reply. Yes, you understood my point correctly. Will add this in our roadmap.

Thanks,
Pierre