meng2003/music2dance

1import torch.nn as nn2import torch.nn.functional as F3import torch4import numpy as np5import scipy.linalg6 7class InvertibleConv1x1(nn.Module):8    def __init__(self, num_channels, LU_decomposed=True):9        super().__init__()10        w_shape = [num_channels, num_channels]11        w_init = np.linalg.qr(np.random.randn(*w_shape))[0].astype(np.float32)12        if not LU_decomposed:13            # Sample a random orthogonal matrix:14            self.register_parameter("weight", nn.Parameter(torch.Tensor(w_init)))15        else:16            # import pdb;pdb.set_trace()17            np_p, np_l, np_u = scipy.linalg.lu(w_init)18            np_s = np.diag(np_u)19            np_sign_s = np.sign(np_s)20            np_log_s = np.log(np.abs(np_s))21            np_u = np.triu(np_u, k=1)22            l_mask = np.tril(np.ones(w_shape, dtype=np.float32), -1)23            eye = np.eye(*w_shape, dtype=np.float32)24            self.register_buffer('p', torch.Tensor(np_p.astype(np.float32)))25            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(np.float32)))26            self.l = nn.Parameter(torch.Tensor(np_l.astype(np.float32)))27            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(np.float32)))28            self.u = nn.Parameter(torch.Tensor(np_u.astype(np.float32)))29            self.l_mask = torch.Tensor(l_mask)30            self.eye = torch.Tensor(eye)31        self.w_shape = w_shape32        self.LU = LU_decomposed33        self.first_pass = True34        self.saved_weight = None35        self.saved_dsldj = None36 37    def get_weight(self, input, reverse):38        w_shape = self.w_shape39        if not self.LU:40            dlogdet = torch.slogdet(self.weight)[1] * input.size(2) * input.size(3)41            if not reverse:42                weight = self.weight.view(w_shape[0], w_shape[1], 1, 1)43            else:44                weight = torch.inverse(self.weight.double()).float()\45                              .view(w_shape[0], w_shape[1], 1, 1)46            return weight, dlogdet47        else:48            self.p = self.p.to(input.device)49            self.sign_s = self.sign_s.to(input.device)50            self.l_mask = self.l_mask.to(input.device)51            self.eye = self.eye.to(input.device)52            l = self.l * self.l_mask + self.eye53            u = self.u * self.l_mask.transpose(0, 1).contiguous() + torch.diag(self.sign_s * torch.exp(self.log_s))54            dlogdet = self.log_s.sum() * input.size(2) * input.size(3)55            if not reverse:56                w = torch.matmul(self.p, torch.matmul(l, u))57            else:58                l = torch.inverse(l.double()).float()59                u = torch.inverse(u.double()).float()60                w = torch.matmul(u, torch.matmul(l, self.p.inverse()))61            return w.view(w_shape[0], w_shape[1], 1, 1), dlogdet62 63    def forward(self, x, cond, sldj=None, reverse=False):64        """65        log-det = log|abs(|W|)| * pixels66        """67        x = torch.cat(x, dim=1)68        if not reverse:69            weight, dsldj = self.get_weight(x, reverse)70        else:71            if self.first_pass:72                weight, dsldj = self.get_weight(x, reverse)73                self.saved_weight = weight74                if sldj is not None:75                    self.saved_dsldj = dsldj76                self.first_pass = False77            else:78                weight = self.saved_weight79                if sldj is not None:80                    dsldj = self.saved_dsldj81 82        if not reverse:83            x = F.conv2d(x, weight)84            if sldj is not None:85                sldj = sldj + dsldj86        else:87            x = F.conv2d(x, weight)88            if sldj is not None:89                sldj = sldj - dsldj90        x = x.chunk(2, dim=1)91        return x, sldj92 93 94class InvConv(nn.Module):95    """Invertible 1x1 Convolution for 2D inputs. Originally described in Glow96    (https://arxiv.org/abs/1807.03039). Does not support LU-decomposed version.97 98    Args:99        num_channels (int): Number of channels in the input and output.100        random_init (bool): Initialize with a random orthogonal matrix.101            Otherwise initialize with noisy identity.102    """103    def __init__(self, num_channels, random_init=False):104        super(InvConv, self).__init__()105        self.num_channels = num_channels106 107        if random_init:108            # Initialize with a random orthogonal matrix109            w_init = np.random.randn(self.num_channels, self.num_channels)110            w_init = np.linalg.qr(w_init)[0]111        else:112            # Initialize as identity permutation with some noise113            w_init = np.eye(self.num_channels, self.num_channels) \114                     + 1e-3 * np.random.randn(self.num_channels, self.num_channels)115        self.weight = nn.Parameter(torch.from_numpy(w_init.astype(np.float32)))116 117    def forward(self, x, cond, sldj, reverse=False):118        x = torch.cat(x, dim=1)119 120        ldj = torch.slogdet(self.weight)[1] * x.size(2) * x.size(3)121 122        if reverse:123            weight = torch.inverse(self.weight.double()).float()124            sldj = sldj - ldj125        else:126            weight = self.weight127            sldj = sldj + ldj128 129        weight = weight.view(self.num_channels, self.num_channels, 1, 1)130        x = F.conv2d(x, weight)131        x = x.chunk(2, dim=1)132 133        return x, sldj134