scheduling a loss weight parameter with PTL

[davidegraff]

I'm currently training a VAE with PTL, so my loss is the weighted sum of two losses: a reconstruction term and a KL term. That is, a training step looks (roughly) like this:

def training_step(self, batch: Tuple, batch_idx) -> torch.Tensor:
    X, _ = batch

    Z_mean, Z_logvar = self.encode(X)
    Z = self.sample(Z_mean, Z_logvar)
    X_hat = self.decode(Z)

    L_rec = F.binary_cross_entropy(X_hat, X, reduction="sum")
    L_kl = -0.5 * (1 + Z_logvar - Z_mean**2 - Z_logvar.exp()).sum()

    L_elbo = L_rec + self.v_kl * L_kl
    return L_elbo

Because the weighting applied to the KL term, self.v_kl in this snippet, can drastically impact model training and performance, I'm looking at ways to optimize this. According to this this blog post, a cyclical schedule for the weight of the KL term is a promising strategy. The original code used native pytorch and manually adjusted the weight for each training sample (c.f., L178 and L203.)

In it's simplest form, I would like v_kl to monotonically increase from 0 to 1 throughout the course of an epoch (i.e., frequency="step"). I can do this manually by adding the line:
v_kl = batch_idx / (self.trainer.estimated_stepping_batches // self.trainer.max_epochs)
but I'm wondering if it would be possible to "automate" this? My reasoning is that I might want to modularize the scheduling of the weight later on using the different modes of the CyclicLR scheduler. Doing so right now would require "reinventing the wheel" of this LR scheduler, and I'd like to avoid that if possible. Does anyone have any thoughts? Thanks for the help!