Custom node type mismatch: expected type: <class 'flax.core.frozen_dict.FrozenDict'>, value: Traced<ShapedArray(float32[256,14,14,3])>with<DynamicJaxprTrace(level=0/1)>.

[mwitiderrick]

Trying to implement batch norm, Am I doing something wrong?

class CNN(nn.Module):

  @nn.compact
  def __call__(self, x, training):
    x = nn.Conv(features=32, kernel_size=(3, 3))(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.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  
    x = nn.Dense(features=256)(x)
    x = nn.BatchNorm(use_running_average=not training)(x)
    x = nn.Dropout(0.2, deterministic=not training)(x)
    x = nn.relu(x)
    x = nn.Dense(features=2)(x)
    x = nn.log_softmax(x)
    return x

ValueError                                Traceback (most recent call last)
Input In [28], in <cell line: 6>()
      5 testing_accuracy = []
      6 for epoch in range(1, num_epochs + 1):
----> 7     train_state, train_metrics = train_one_epoch(state, train_loader)
      8     training_loss.append(train_metrics['loss'])
      9     training_accuracy.append(train_metrics['accuracy'])

Input In [21], in train_one_epoch(state, dataloader)
      4 for cnt, (images, labels) in enumerate(dataloader):
      5     images = images / 255.0
----> 6     state, metrics = train_step(state, images, labels)
      7     batch_metrics.append(metrics)
      9 batch_metrics_np = jax.device_get(batch_metrics)  

    [... skipping hidden 14 frame]

Input In [19], in train_step(state, images, labels)
      4 rng, subrng = jax.random.split(state.rng)
      5 (batch_loss, (logits, batch_stats)), grads= jax.value_and_grad(compute_loss, has_aux=True)(images,labels,state.batch_stats,state.params)
----> 6 state = state.apply_gradients(grads=grads) 
      7 # state = state.update_batch_stats(state, batch_stats['batch_stats'])
      8 # state = state.update_rng(state, rng)
     10 metrics = compute_metrics(logits=logits, labels=labels) 

File ~/miniforge3/envs/layer38/lib/python3.8/site-packages/flax/training/train_state.py:75, in TrainState.apply_gradients(self, grads, **kwargs)
     59 """Updates `step`, `params`, `opt_state` and `**kwargs` in return value.
     60 
     61 Note that internally this function calls `.tx.update()` followed by a call
   (...)
     71   replaced as specified by `kwargs`.
     72 """
     73 updates, new_opt_state = self.tx.update(
     74     grads, self.opt_state, self.params)
---> 75 new_params = optax.apply_updates(self.params, updates)
     76 return self.replace(
     77     step=self.step + 1,
     78     params=new_params,
     79     opt_state=new_opt_state,
     80     **kwargs,
     81 )

File ~/miniforge3/envs/layer38/lib/python3.8/site-packages/optax/_src/update.py:42, in apply_updates(params, updates)
     24 def apply_updates(params: base.Params, updates: base.Updates) -> base.Params:
     25   """Applies an update to the corresponding parameters.
     26 
     27   This is a utility functions that applies an update to a set of parameters, and
   (...)
     40     Updated parameters, with same structure, shape and type as `params`.
     41   """
---> 42   return jax.tree_util.tree_map(
     43       lambda p, u: jnp.asarray(p + u).astype(jnp.asarray(p).dtype),
     44       params, updates)

    [... skipping hidden 1 frame]

File ~/miniforge3/envs/layer38/lib/python3.8/site-packages/jax/_src/tree_util.py:204, in <listcomp>(.0)
    171 """Maps a multi-input function over pytree args to produce a new pytree.
    172 
    173 Args:
   (...)
    201   [[5, 7, 9], [6, 1, 2]]
    202 """
    203 leaves, treedef = tree_flatten(tree, is_leaf)
--> 204 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
    205 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))

ValueError: Custom node type mismatch: expected type: <class 'flax.core.frozen_dict.FrozenDict'>, value: Traced<ShapedArray(float32[256,14,14,3])>with<DynamicJaxprTrace(level=0/1)>.

[jheek]

jax.value_and_grad computes the loss wrt the first argument. So in this code snippet you are computing the gradients wrt the inputs instead of the params. You can fix this by passing the model params as the first argument of the loss function

[mwitiderrick]

@jheek, thanks for the response. You are right!