Using precomputed mean and var for BatchNorm

[matthias-wright]

Hi all,

I am trying to load precomputed batch stats into the flax.nn.normalization.BatchNorm module.

From the source code I gather that both the mean and the variance are registered as a state.

Here is what I tried:

import flax.nn as nn
from jax import random

key = random.PRNGKey(0)

scale = random.normal(key, shape=(3,))
bias = random.normal(key, shape=(3,))
mean = random.normal(key, shape=(3,))
var = random.normal(key, shape=(3,))

x = random.normal(key, shape=(1, 5, 5, 3))

batch_stats = nn.Collection({'mean': mean, 'var': var}).mutate()
bn = nn.normalization.BatchNorm.partial(bias_init=lambda *_ : bias, 
                                        scale_init=lambda *_ : scale, 
                                        use_running_average=True, 
                                        batch_stats=batch_stats)

x_out = bn.init(key, x)

The error message I get is:
AttributeError: '_GeneratorContextManager' object has no attribute 'retrieve'

Has anyone here tried this before?

Thanks!

[avital]

Hi @matthias-wright -- you're using the deprecated flax.nn API -- please try using flax.linen instead, where "batch_stats" are simply another variable collection just like "params".

[avital]

(Keep us posted if you can't figure it out and we can share a code snippet, but I hope it should be pretty straightforward)

[matthias-wright]

Hi @avital, thank you for the answer!
In the new linen API, "batch_stats" is no longer a parameter for the BatchNorm module.
How can I pass this collection to a BatchNorm module?

[avital]

Hi @matthias-wright -- I wrote up a Colab with sample code that shows how to do this. Note that module instances at the top level of your training code in Flax/Linen don't actually hold their variables. Variables are passed in when you call apply. This is to ensure that Flax interacts with JAX in a purely functional way. So you don't construct a batch norm instance with variables, you modify your variable dictionary directly.

Here is the Colab: https://colab.research.google.com/drive/13ltYm14azyWVYUYUsR_UIHQqlG25uo_M?usp=sharing

And code below:

import flax.linen as nn
from jax import random
from flax.core import freeze, unfreeze

key = random.PRNGKey(0)

scale = random.normal(key, shape=(3,))
bias = random.normal(key, shape=(3,))
mean = random.normal(key, shape=(3,))
var = random.normal(key, shape=(3,))

x = random.normal(key, shape=(1, 5, 5, 3))

bn = nn.normalization.BatchNorm(bias_init=lambda *_ : bias, 
                                scale_init=lambda *_ : scale, 
                                use_running_average=True)

initialized_variables = bn.init(key, x)
print("initialized variables: ")
print(initialized_variables)

# NOTE TO SELF: Maybe we should not return FrozenDicts from `init`?
variables_with_preloaded_batch_stats = unfreeze(initialized_variables)
variables_with_preloaded_batch_stats['batch_stats']['mean'] = mean
variables_with_preloaded_batch_stats['batch_stats']['var'] = var
variables_with_preloaded_batch_stats = freeze(variables_with_preloaded_batch_stats)

print()
print("variable with preloaded batch stats: ")
print(variables_with_preloaded_batch_stats)

output = bn.apply(variables_with_preloaded_batch_stats, x)

[matthias-wright]

Hi @avital, thank you very much!