[CLOSED] Importance Weighted Autoencoder about scvi-tools HOT 1 CLOSED

    def iwae_loss(self, ys, k=10):
        tall_shape = [ys.size(0) * k] + list(ys.size()[1:])
        ys_tall = ys.unsqueeze(1).repeat(1, k, 1, 1).view(tall_shape)

        wide_shape = [ys.size(0), k]
        vae_losses = self.vae_losses(ys_tall).view(wide_shape)

        iwae_gains = logsumexp(-1 * vae_losses, dim=1) - np.log(k)
        return -1 * iwae_gains.sum()

def logsumexp(inputs, dim=None, keepdim=False):
    return (inputs - f.log_softmax(inputs, dim=dim)).mean(dim, keepdim=keepdim)