meng2003/music2dance

1import numpy as np2import torch.nn as nn3import torch.nn.utils as utils4 5 6def bits_per_dim(x, nll):7    """Get the bits per dimension implied by using model with `loss`8    for compressing `x`, assuming each entry can take on `k` discrete values.9 10    Args:11        x (torch.Tensor): Input to the model. Just used for dimensions.12        nll (torch.Tensor): Scalar negative log-likelihood loss tensor.13 14    Returns:15        bpd (torch.Tensor): Bits per dimension implied if compressing `x`.16    """17    dim = np.prod(x.size()[1:])18    bpd = nll / (np.log(2) * dim)19 20    return bpd21 22 23def clip_grad_norm(optimizer, max_norm, norm_type=2):24    """Clip the norm of the gradients for all parameters under `optimizer`.25 26    Args:27        optimizer (torch.optim.Optimizer):28        max_norm (float): The maximum allowable norm of gradients.29        norm_type (int): The type of norm to use in computing gradient norms.30    """31    for group in optimizer.param_groups:32        utils.clip_grad_norm_(group['params'], max_norm, norm_type)33 34 35class NLLLoss(nn.Module):36    """Negative log-likelihood loss assuming isotropic gaussian with unit norm.37 38    Args:39        k (int or float): Number of discrete values in each input dimension.40            E.g., `k` is 256 for natural images.41 42    See Also:43        Equation (3) in the RealNVP paper: https://arxiv.org/abs/1605.0880344    """45    def __init__(self, k=256):46        super(NLLLoss, self).__init__()47        self.k = k48 49    def forward(self, z, sldj):50        prior_ll = -0.5 * (z ** 2 + np.log(2 * np.pi))51        prior_ll = prior_ll.flatten(1).sum(-1) \52            - np.log(self.k) * np.prod(z.size()[1:])53        ll = prior_ll + sldj54        nll = -ll.mean()55 56        return nll57