Optimization groups

[aegroto]

Hello! What's the most idiomatic way to have optimization groups in Burn? Although I wouldn't se different optimizers, I need to use different learning rates for parameters in the same module.

[nathanielsimard]

We dont support optimization groups just yet, at least with a convenient API. However, you could create multiple optimizers with different learning rates and extract the gradients from the backward pass into multiple gradients params. You can do that using a visitor by grouping the gradients into each parameter group you want.

[aegroto]

Thanks! I will leave here a snippet for reference, in case anyone has my same need, although it's a pretty rough way to achieve separate optimization groups as I am not using a visitor:

// Calculate grads through backward pass for all modules
let mut grads = loss.backward();

// Extract params only for the "seed" module by using grad_remove
let mut seed_grads_params = GradientsParams::new();

let seed_grads = model.seed.current_latent.features.grad_remove(&mut grads).unwrap();
// Note that the ParamId should be known as there is no way as of now to access it from the Param object
seed_grads_params.register(ParamId::from("current_latent"), seed_grads);
model = seed_optimizer.step(
    self.state.seed_lr,
    model,
    seed_grads_params,
);

// Optimize the remaining modules
let grads_params = GradientsParams::from_grads(grads, &model);
model = optimizer.step(self.state.lr, model, grads_params);