Is there a way to sum the parameters of flax.nn.Module and create a new module?

[BoyuanJackChen]

I am training a flax Adam optimizer. The data is divided batches, and I want to update the model with the averaged grad among all batches. I checked that the type of grad is flax.nn.Module, and I wonder how to sum it up with jax.

Description of the model to be implemented

The model is a simple one with 8 layers:

class NeRF_Model(nn.Module):
    def apply(self, x):
        the_input = x
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = jnp.append(arr=x, values=the_input, axis=1)  # Dense_4: (257, 256)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=256)
        x = nn.relu(x)
        x = nn.Dense(x, features=4, kernel_init = initializers.normal(1),
                     bias_init=initializers.ones)
        # output: (r,g,b, \upsigma)
        return x

The update line looks like this:

        def loss_fn(model):  # Will be used repeatedly in the loop
            rgb, depth, acc = render_rays(model, rays_o, rays_d, near=near, far=far,
                                batchify_size=batchify_size, N_samples=N_samples, rand=True)
            loss = jnp.mean(jnp.square(rgb - target))
            return loss, rgb
        grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
        (loss, logits), grad = grad_fn(optimizer.target)
        optimizer = optimizer.apply_gradient(grad)

Dataset the model could be trained on

Image data.

Specific points to consider

Reference implementations in other frameworks

[jheek]

Hi,
if I understand correctly you have a loop over multiple batches and you want to average the gradients across the batches.

Here's a code snippet that should give the basic idea of how to do this:

def loss_fn(model, batch):  # Will be used repeatedly in the loop
            rays_o, rays_d, target = batch # or whatever is in the batch
            rgb, depth, acc = render_rays(model, rays_o, rays_d, near=near, far=far,
                                batchify_size=batchify_size, N_samples=N_samples, rand=True)
            loss = jnp.mean(jnp.square(rgb - target))
            return loss, rgb
grad_fn = jax.value_and_grad(loss_fn, has_aux=True)

avg_grad = None
num_batches = len(batches)
for batch in batches:
  (loss, logits), grad = grad_fn(optimizer.target, batch)
  if avg_grad is None:
    avg_grad = jax.tree_map(lambda g: g / num_batches, grad)
  else:
    avg_grad = jax.tree_multimap(lambda avg, g: avg + g / num_batches, avg_grad, grad)

optimizer = optimizer.apply_gradient(avg_grad)

[BoyuanJackChen]

Thanks! This is very helpful!