How to use nn.BatchNorm?

[fmscole]

My code is:

class CNN(nn.Module):
  """A simple CNN model."""

  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=True)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.BatchNorm()(x,use_running_average=True)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x

But ,get the errors:

~/anaconda3/lib/python3.8/site-packages/flax/core/scope.py in variable(self, col, name, init_fn, *init_args)
    716       if not self.is_mutable_collection(col) or init_fn is None:
    717         if self.is_collection_empty(col):
--> 718           raise errors.ScopeCollectionNotFound(col, name, self.path_text)
    719         raise errors.ScopeVariableNotFoundError(name, col, self.path_text)
    720       init_value = init_fn(*init_args)

ScopeCollectionNotFound: Tried to access "mean" from collection "batch_stats" in "/BatchNorm_0" but the collection is empty. (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeCollectionNotFound)

[cgarciae]

Update: for a more indepth explanation please take a look at our Batch normalization guide.

---

Hey @fmscole, to use BatchNorm you usually follow a pattern like this:

class CNN(nn.Module):
  """A simple CNN model."""
  @nn.compact
  def __call__(self, x, training: bool):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x
    
x = jnp.ones((16, 28, 28, 1))
module = CNN()

# initialize weights
key = jax.random.PRNGKey(0)
variables = module.init(key, x, training=False)

# forward during training
y, updates = module.apply(variables, x, training=True, mutable=['batch_stats'])
variables = variables.copy(updates)

# foward during testing/inference
y = module.apply(variables, x, training=False)

[fmscole]

@cgarciae Thanks for your reply!

[zaccharieramzi]

@cgarciae thanks for your answer !
I wanted to know if it was possible to use BatchNorm in a way that does not store the running statistics but rather computes the mean and variance on the incoming batch? The reason I want to do this is because it's done in MAML (see here).

If I just do something like x = fnn.BatchNorm(use_running_average=False)(x), I get the error from the OP.

[zaccharieramzi]

Got it, I just need to pass in the batch statistics no matter what

[fmscole]

When I use the batch_stats that retruned from loss function:

@jax.jit
def predict(state, batch_stats,image_i):
  logits = state.apply_fn({'params': state.params,"batch_stats":batch_stats},image_i,is_training=False)
  return logits

result is:

So I abandon the batch_stats ：

@jax.jit
def predict(state, batch_stats,image_i):
  logits,_ = state.apply_fn({'params': state.params},image_i,is_training=True, mutable=['batch_stats'])
  return logits

result is:

Is this normal?
Why？

[fmscole]

I already know how to use batchnormal:

from flax import linen as nn
from flax.metrics import tensorboard
from flax.training import train_state
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import optax
from sqlalchemy import false
import tensorflow_datasets as tfds


class CNN(nn.Module):
  @nn.compact
  def __call__(self, x,is_training:bool=True):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not is_training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=not is_training)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x



@jax.jit
def apply_model(state, images, labels,old_batch_stats):
  def loss_fn(params,old_batch_stats):
    logits,mutated_vars = state.apply_fn({'params': params,"batch_stats":old_batch_stats}, images,is_training=True, mutable=['batch_stats'])
    one_hot = jax.nn.one_hot(labels, 10)
    loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
    return loss, (logits,mutated_vars['batch_stats'])
    
  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (loss, (logits,new_batch_stats)), grads = grad_fn(state.params,old_batch_stats)
  accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
  return grads, loss, accuracy,new_batch_stats


@jax.jit
def update_model(state, grads):
  return state.apply_gradients(grads=grads)


def train_epoch(state, train_ds, batch_size, rng,batch_stats):
  train_ds_size = len(train_ds['image'])
  steps_per_epoch = train_ds_size // batch_size

  perms = jax.random.permutation(rng, len(train_ds['image']))
  perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
  perms = perms.reshape((steps_per_epoch, batch_size))

  epoch_loss = []
  epoch_accuracy = []
  for perm in perms:
    batch_images = train_ds['image'][perm, ...]
    batch_labels = train_ds['label'][perm, ...]
    grads, loss, accuracy ,batch_stats= apply_model(state, batch_images, batch_labels,batch_stats)
    
    state = update_model(state, grads)
    
    epoch_loss.append(loss)
    epoch_accuracy.append(accuracy)
  train_loss = np.mean(epoch_loss)
  train_accuracy = np.mean(epoch_accuracy)
  return state, train_loss, train_accuracy,batch_stats


def get_datasets():
  ds_builder = tfds.builder('mnist')
  ds_builder.download_and_prepare()
  train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))
  test_ds = tfds.as_numpy(ds_builder.as_dataset(split='test', batch_size=-1))
  train_ds['image'] = jnp.float32(train_ds['image']) / 255.
  test_ds['image'] = jnp.float32(test_ds['image']) / 255.
  return train_ds, test_ds


def create_train_state(rng):
  cnn = CNN()
  variables=cnn.init(rng, jnp.ones([1, 28, 28, 1]))
  params = variables['params']
  batch_stats=variables['batch_stats']
  tx = optax.sgd(0.01, 0.99)
  state=train_state.TrainState.create(apply_fn=cnn.apply, params=params, tx=tx)
  return state,batch_stats

@jax.jit
def predict(state, batch_stats,image_i):
  logits= state.apply_fn({'params': state.params,"batch_stats":batch_stats},image_i,is_training=False)
  return logits
def test(state, batch_stats,test_ds):
  images = test_ds['image']
  labels = test_ds['label']
  batchs=1000
  accuracy=0
  for i in range(0,len(images),batchs):
    image_i=images[i:i+batchs]
    label_i=labels[i:i+batchs]

    logits= predict(state, batch_stats,image_i)
    accuracy += jnp.sum(jnp.argmax(logits, -1) == label_i)
  
  return accuracy/len(images)

def train_and_evaluate() -> train_state.TrainState:
  train_ds, test_ds = get_datasets()
  rng = jax.random.PRNGKey(0)

  
  rng, init_rng = jax.random.split(rng)
  state,batch_stats = create_train_state(init_rng)
  
  for epoch in range(1, 100 + 1):
    rng, input_rng = jax.random.split(rng)
    state, train_loss, train_accuracy,batch_stats = train_epoch(state, train_ds,
                                                    100,
                                                    input_rng,batch_stats)
    
    
    print(test(state, batch_stats,test_ds),end=" ")
  
  return state

train_and_evaluate() 

[fmscole]

or like this:

from flax import linen as nn
from flax.metrics import tensorboard
from flax.training import train_state
import jax
import jax.numpy as jnp
import ml_collections
import numpy as np
import optax
import tensorflow_datasets as tfds

class CNN(nn.Module):
  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=False)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.BatchNorm(use_running_average=False)(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x

@jax.jit
def apply_model(state, images, labels):
  def loss_fn(params):
    logits,batch_stats = state.apply_fn({'params': params}, images, mutable=['batch_stats'])
    one_hot = jax.nn.one_hot(labels, 10)
    loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
    return loss, logits
  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (loss, logits), grads = grad_fn(state.params)
  accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
  return grads, loss, accuracy

@jax.jit
def update_model(state, grads):
  return state.apply_gradients(grads=grads)

def train_epoch(state, train_ds, batch_size, rng):
  train_ds_size = len(train_ds['image'])
  steps_per_epoch = train_ds_size // batch_size

  perms = jax.random.permutation(rng, len(train_ds['image']))
  perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
  perms = perms.reshape((steps_per_epoch, batch_size))

  epoch_loss = []
  epoch_accuracy = []

  for perm in perms:
    batch_images = train_ds['image'][perm, ...]
    batch_labels = train_ds['label'][perm, ...]
    grads, loss, accuracy = apply_model(state, batch_images, batch_labels)
    state = update_model(state, grads)
    
    epoch_loss.append(loss)
    epoch_accuracy.append(accuracy)
  train_loss = np.mean(epoch_loss)
  train_accuracy = np.mean(epoch_accuracy)
  return state, train_loss, train_accuracy


def get_datasets():
  ds_builder = tfds.builder('mnist')
  ds_builder.download_and_prepare()
  train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))
  test_ds = tfds.as_numpy(ds_builder.as_dataset(split='test', batch_size=-1))
  train_ds['image'] = jnp.float32(train_ds['image']) / 255.
  test_ds['image'] = jnp.float32(test_ds['image']) / 255.
  return train_ds, test_ds


def create_train_state(rng):
  cnn = CNN()
  variables=cnn.init(rng, jnp.ones([1, 28, 28, 1]))
  params = variables['params']
  tx = optax.sgd(0.01, 0.99)
  return train_state.TrainState.create(
      apply_fn=cnn.apply, params=params, tx=tx)

@jax.jit
def predict(state,image_i):
  logits,batch_stats = state.apply_fn({'params': state.params},image_i, mutable=['batch_stats'])
  return logits
def test(state,test_ds):
  images = test_ds['image']
  labels = test_ds['label']
  batchs=1000
  accuracy=0
  for i in range(0,len(images),batchs):
    image_i=images[i:i+batchs]
    label_i=labels[i:i+batchs]

    logits= predict(state,image_i)
    accuracy += jnp.sum(jnp.argmax(logits, -1) == label_i)
  
  return accuracy/len(images)

def train_and_evaluate() -> train_state.TrainState:
  train_ds, test_ds = get_datasets()
  rng = jax.random.PRNGKey(0)

  
  rng, init_rng = jax.random.split(rng)
  state = create_train_state(init_rng)
  
  for epoch in range(1, 100 + 1):
    rng, input_rng = jax.random.split(rng)
    state, train_loss, train_accuracy = train_epoch(state, train_ds,
                                                    100,
                                                    input_rng)
    print(test(state, test_ds),end=" ")
    
  return state

train_and_evaluate()