patrick-kidger
commented
2 years ago Right. So this is something I've been avoiding implementing. It is possible using Equinox but I'm honestly not sure it belongs in the library itself, as it requires changing how you use things slightly.
First of all, an (untested) implementation looks something like this.
def pvar(x, axis_name):
mean = lax.pmean(x, axis_name)
var = lax.pmean((x - mean)**2, axis_name)
return var
class BatchNorm(eqx.Module):
running_mean: Optional[jnp.ndarray]
running_var: Optional[jnp.ndarray]
momentum: float
axis_name: str
eps: float
update_stats: bool
def __init__(self, momentum=0.99, axis_name="batch", eps=1e-5, update_stats=True, **kwargs):
super().__init__(**kwargs)
self.running_mean = None
self.running_var = None
self.momentum = momentum
self.axis_name = axis_name
self.eps = eps
self.update_stats = update_stats
def __call__(self, x, update_stats=None):
if update_stats is None:
update_stats = self.update_stats
mean = lax.pmean(x, axis=self.axis_name)
var = pvar(x, axis=self.axis_name)
if self.running_mean is None:
running_mean = mean
running_var = var
else:
running_mean = self.running_mean
running_var = self.running_var
running_mean = (1 - self.momentum) * mean + self.momentum * lax.stop_gradient(running_mean)
running_var = (1 - self.momentum) * var + self.momentum * lax.stop_gradient(running_var)
x = (x - running_mean) / jnp.sqrt(running_var + self.eps)
if self.update_stats:
object.__setattr__(self, "running_mean", running_mean)
object.__setattr__(self, "running_var", running_var)
return xThe really weird thing here is the use of object.__setattr__. This is essentially equivalent to self.running_mean = running_mean; self.running_var = running_var, but hackily avoids the error message you get about trying to mutate something. (eqx.Modules are supposed to be immutable after __init__.)
Assuming you've wrapped everything up into a JIT wrapper then this means you will need to pass out anything you've mutated. That is, your loss computation needs to change from something looking like
@jax.jit
@jax.grad
def loss(model, x, y):
pred_y = jax.vmap(model)(x)
return jnp.mean((y - pred_y)**2)to something looking like
@jax.jit
@jax.grad
def loss(model, x, y):
pred_y = jax.vmap(model, axis_name="batch")(x)
return jnp.mean((y - pred_y)**2), modeland you'll need to use the updated version of model -- which has the updated batch statistics on it.
Forgetting to use the updated version of model is a potential footgun, which is why I'm a bit concerned about having it in the library. (Flax gets around this by just having footguns all over the place, so having another new one isn't a surprise...)
Other minor things to note.
- The use of named axes to handle the fact that we're wrapping in
vmap. - The use of
lax.stop_gradient. We never want to update these arrays via gradient descent. - The
update_statsflag, used to disable updates during inference. (Neither of those are particularly weird, at least.)
marcelroed
In Flax, Batch Normalization is a bit finicky since each call to
applyrequires markingbatch_statsas mutable and updating thebatch_statsafterward.How could this be implemented as a Module in Equinox? I'm happy to submit an implementation given some guidance.