meng2003/music2dance

1import torch2import torch.nn as nn3import torch.nn.functional as F4 5 6class Flip(nn.Module):7    def forward(self, x, cond, sldj, reverse=False):8        assert isinstance(x, tuple) and len(x) == 29        return (x[1], x[0]), sldj10 11 12def mean_dim(tensor, dim=None, keepdims=False):13    """Take the mean along multiple dimensions.14 15    Args:16        tensor (torch.Tensor): Tensor of values to average.17        dim (list): List of dimensions along which to take the mean.18        keepdims (bool): Keep dimensions rather than squeezing.19 20    Returns:21        mean (torch.Tensor): New tensor of mean value(s).22    """23    if dim is None:24        return tensor.mean()25    else:26        if isinstance(dim, int):27            dim = [dim]28        dim = sorted(dim)29        for d in dim:30            tensor = tensor.mean(dim=d, keepdim=True)31        if not keepdims:32            for i, d in enumerate(dim):33                tensor.squeeze_(d-i)34        return tensor35 36 37def checkerboard(x, reverse=False):38    """Split x in a checkerboard pattern. Collapse horizontally."""39    # Get dimensions40    if reverse:41        b, c, h, w = x[0].size()42        w *= 243        device = x[0].device44    else:45        b, c, h, w = x.size()46        device = x.device47 48    # Get list of indices in alternating checkerboard pattern49    y_idx = []50    z_idx = []51    for i in range(h):52        for j in range(w):53            if (i % 2) == (j % 2):54                y_idx.append(i * w + j)55            else:56                z_idx.append(i * w + j)57    y_idx = torch.tensor(y_idx, dtype=torch.int64, device=device)58    z_idx = torch.tensor(z_idx, dtype=torch.int64, device=device)59 60    if reverse:61        y, z = (t.contiguous().view(b, c, h // 2 * w) for t in x)62        x = torch.zeros(b, c, h * w, dtype=y.dtype, device=y.device)63        x[:, :, y_idx] += y64        x[:, :, z_idx] += z65        x = x.view(b, c, h, w)66 67        return x68    else:69        if h % 2 != 0:70            raise RuntimeError('Checkerboard got odd height input: {}'.format(h))71 72        x = x.view(b, c, h * w)73        y = x[:, :, y_idx].view(b, c, h // 2, w)74        z = x[:, :, z_idx].view(b, c, h // 2, w)75 76        return y, z77 78 79def channelwise(x, reverse=False):80    """Split x channel-wise."""81    if reverse:82        x = torch.cat(x, dim=1)83        return x84    else:85        y, z = x.chunk(2, dim=1)86        return y, z87 88 89def squeeze(x):90    """Trade spatial extent for channels. I.e., convert each91    1x4x4 volume of input into a 4x1x1 volume of output.92 93    Args:94        x (torch.Tensor): Input to squeeze.95 96    Returns:97        x (torch.Tensor): Squeezed or unsqueezed tensor.98    """99    # import pdb; pdb.set_trace()100    b, c, h, w = x.size()101    x = x.view(b, c, h // 2, 2, w, 1)102    x = x.permute(0, 1, 3, 5, 2, 4).contiguous()103    x = x.view(b, c * 2, h // 2, w)104 105    return x106 107 108def unsqueeze(x):109    """Trade channels channels for spatial extent. I.e., convert each110    4x1x1 volume of input into a 1x4x4 volume of output.111 112    Args:113        x (torch.Tensor): Input to unsqueeze.114 115    Returns:116        x (torch.Tensor): Unsqueezed tensor.117    """118    b, c, h, w = x.size()119    x = x.view(b, c // 2, 2, 1, h, w)120    x = x.permute(0, 1, 4, 2, 5, 3).contiguous()121    x = x.view(b, c // 2, h * 2, w)122 123    return x124 125 126def concat_elu(x):127    """Concatenated ReLU (http://arxiv.org/abs/1603.05201), but with ELU."""128    return F.elu(torch.cat((x, -x), dim=1))129 130 131def safe_log(x):132    return torch.log(x.clamp(min=1e-22))133