I use lenet5 for mnist classification. but the network performance is extremely bad, I notice some people say KL annealing would be helpful, to scale the KL term in the ELBO, one potential method is using poutine.scale as following
SVI(poutine.scale(model, scalar),poutine.scale( guide, scalar), opt, loss=Trace_ELBO())
the scalar is the scale parameter. during training, we use svi.step()
but how can we change the scalar value during training?