Optimizer Step Closure with Sampling

[Leon0402]

Hi,

I saw the the IVON Optimizer supports passing a closure to the step function: https://github.com/team-approx-bayes/ivon/blob/main/ivon/_ivon.py#L170

I was wondering whether it would make sense to add the sampling to this function, i.e.:

 losses = []
for _ in range(self.mc_samples):
    with torch.enable_grad():
        with self.sampled_params(train=True):
            loss = closure()
    losses.append(loss)
loss = sum(losses) / self.mc_samples

Now code using the closure syntax would be completely equivalent to the ADAM Implemention, which would help sell the point. This would most drastically be seen in frameworks such as pytorch lightning, which provide abstractions over the training loop (i.e. do things like optimizer step automatically). With this change, IVON should be a drop-in replacement in pytorch lightning for ADAM with zero changes.

@override
def training_step(self, batch, batch_idx) -> torch.FloatTensor:
    output = self.model(batch)
    self.log("train/loss", output.loss)
    return output.loss

Just works :)

With manual optimization mode (e.g. doing this like optimizer.step()) yourself, this becomes much more complicated. I found a way to almost get it by overriding the optimizer_step:

@override
def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_closure):
    def sampled_closure():
        # The optimizer closure will contain our training step, zero grad and backward pass
        with self.optimizer.sampled_params(train=True):
            return optimizer_closure()

    optimizer.step(closure=sampled_closure)

If you are not familiar with pytorch lightning, I think this might be the best very short introduction on this specific part:
https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#training-loop

As you can see it abstracts the training loop away. But the information here is not entirely correct, because they actually put:

loss = training_step(batch, batch_idx)

# clear gradients
optimizer.zero_grad()

# backward
loss.backward()

in a closure and pass this to optimizer.step(). As backwards needs to happen within the optimizer.sampled_params context manager, this closure needs to be overwritten as can be seen above.