meng2003/music2dance
0
1import math2import torch3import torch.nn as nn4import torch.nn.functional as F5import numpy as np6 7from models.flowplusplus.act_norm import ActNorm, BatchNorm8from models.flowplusplus.inv_conv import InvConv, InvertibleConv1x19from models.flowplusplus.nn import GatedConv10from models.flowplusplus.coupling import Coupling11from models.util import channelwise, checkerboard, Flip, safe_log, squeeze, unsqueeze12 13from models.moglow.modules import GaussianDiag, StudentT14 15class FlowPlusPlus(nn.Module):16 """Flow++ Model17 18 Based on the paper:19 "Flow++: Improving Flow-Based Generative Models20 with Variational Dequantization and Architecture Design"21 by Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, Pieter Abbeel22 (https://openreview.net/forum?id=Hyg74h05tX).23 24 Args:25 scales (tuple or list): Number of each type of coupling layer in each26 scale. Each scale is a 2-tuple of the form27 (num_channelwise, num_checkerboard).28 in_channels (int): Number of channels in the input.29 mid_channels (int): Number of channels in the intermediate layers.30 num_blocks (int): Number of residual blocks in the s and t network of31 `Coupling` layers.32 num_dequant_blocks (int): Number of blocks in the dequantization flows.33 """34 def __init__(self,35 scales=((0, 4), (2, 3)),36 in_shape=(3, 32, 32),37 cond_dim=0,38 mid_channels=96,39 num_blocks=10,40 num_components=32,41 use_attn=True,42 use_logmix=True,43 use_transformer_nn=False,44 use_pos_emb=False,45 use_rel_pos_emb=False,46 num_heads=10,47 drop_prob=0.2,48 norm_layer=None,49 cond_concat_dims=True,50 cond_seq_len=1,51 flow_dist="normal",52 flow_dist_param=50,53 bn_momentum=0.1):54 super(FlowPlusPlus, self).__init__()55 # Register bounds to pre-process images, not learnable56 self.register_buffer('bounds', torch.tensor([0.9], dtype=torch.float32))57 self.flows = _FlowStep(scales=scales,58 in_shape=in_shape,59 cond_dim=cond_dim,60 mid_channels=mid_channels,61 num_blocks=num_blocks,62 num_components=num_components,63 use_attn=use_attn,64 use_logmix=use_logmix,65 use_transformer_nn=use_transformer_nn,66 use_pos_emb=use_pos_emb,67 use_rel_pos_emb=use_rel_pos_emb,68 num_heads=num_heads,69 drop_prob=drop_prob,70 norm_layer=norm_layer,71 cond_concat_dims=cond_concat_dims,72 cond_seq_len=cond_seq_len,73 bn_momentum=bn_momentum)74 if flow_dist == "normal":75 self.distribution = GaussianDiag()76 elif flow_dist == "studentT":77 in_channels, in_height, in_width = in_shape78 self.distribution = StudentT(flow_dist_param, in_channels)79 80 def forward(self, x, cond, reverse=False):81 if cond is not None:82 cond = cond.permute(0,2,1).unsqueeze(3)83 84 if not reverse: 85 if x is not None:86 x = x.permute(0,2,1).unsqueeze(3)87 else:88 c, h, w = self.flows.z_dim()89 # x = 1.0*torch.randn((cond.size(0), c, h, w), dtype=torch.float32).type_as(cond)90 eps_std=1.091 # x = self.distribution.sample((cond.size(0), c, h, w), eps_std, device=cond.device).type_as(cond)92 assert w==193 x = self.distribution.sample((cond.size(0), c, h), eps_std, device=cond.device).type_as(cond)94 x = x.unsqueeze(-1)95 # import pdb;pdb.set_trace()96 97 sldj = torch.zeros(x.size(0), device=x.device)98 x, sldj = self.flows(x, cond, sldj, reverse)99 100 if reverse:101 if x is not None:102 x = x.squeeze(3).permute(0,2,1)103 104 return x, sldj105 106 def loss_generative(self, z, sldj):107 """Negative log-likelihood loss assuming isotropic gaussian with unit norm.108 109 Args:110 k (int or float): Number of discrete values in each input dimension.111 E.g., `k` is 256 for natural images.112 113 See Also:114 Equation (3) in the RealNVP paper: https://arxiv.org/abs/1605.08803115 """116 # print(z)117 # prior_ll = -0.5 * (z ** 2 + np.log(2 * np.pi))118 # prior_ll = prior_ll.flatten(1).sum(-1)# \119 prior_ll = self.distribution.logp(z)120 prior_ll = prior_ll.flatten(1).sum(-1)# \121 # import pdb;pdb.set_trace()122# - np.log(k) * np.prod(z.size()[1:])123 ll = prior_ll + sldj124 # print(sldj.mean())125 # import pdb;pdb.set_trace()126 nll = -ll.mean()/float(np.log(2.) * z.size(2) * z.size(3))127 # nll = -ll.mean()/float(np.log(2.))128 129 return nll130 131class _FlowStep(nn.Module):132 """Recursive builder for a Flow++ model.133 134 Each `_FlowStep` corresponds to a single scale in Flow++.135 The constructor is recursively called to build a full model.136 137 Args:138 scales (tuple): Number of each type of coupling layer in each scale.139 Each scale is a 2-tuple of the form (num_channelwise, num_checkerboard).140 in_channels (int): Number of channels in the input.141 mid_channels (int): Number of channels in the intermediate layers.142 num_blocks (int): Number of residual blocks in the s and t network of143 `Coupling` layers.144 num_components (int): Number of components in the mixture.145 use_attn (bool): Use attention in the coupling layers.146 drop_prob (float): Dropout probability.147 """148 def __init__(self, scales, in_shape, cond_dim, mid_channels, num_blocks, num_components, use_attn, use_logmix, use_transformer_nn, use_pos_emb, use_rel_pos_emb, num_heads, drop_prob, norm_layer, bn_momentum, cond_concat_dims, cond_seq_len):149 super(_FlowStep, self).__init__()150 in_channels, in_height, in_width = in_shape151 num_channelwise, num_checkerboard = scales[0]152 #import pdb;pdb.set_trace()153 channels = []154 for i in range(num_channelwise):155 new_channels = in_channels// 2156 out_channels = in_channels-new_channels157 # print(norm_layer)158 if norm_layer == "batchnorm":159 channels += [BatchNorm(in_channels, bn_momentum)]160 elif norm_layer == "actnorm":161 channels += [ActNorm(in_channels)]162 if cond_concat_dims:163 c_in_channels = new_channels + cond_dim164 seq_length = in_height165 else:166 c_in_channels = new_channels167 seq_length = in_height + cond_seq_len168 channels += [InvertibleConv1x1(in_channels)]169 channels += [Coupling(in_channels=c_in_channels,170 cond_dim=cond_dim,171 out_channels=out_channels,172 mid_channels=mid_channels,173 num_blocks=num_blocks,174 num_components=num_components,175 use_attn=use_attn,176 use_logmix=use_logmix,177 use_transformer_nn=use_transformer_nn,178 use_pos_emb=use_pos_emb,179 use_rel_pos_emb=use_rel_pos_emb,180 num_heads=num_heads,181 seq_length=seq_length,182 output_length=in_height,183 concat_dims=cond_concat_dims,184 drop_prob=drop_prob)]#,185 #Flip()] Flip currently does not work with odd number of channels. But is it needed when we have channel mixing with 1x1convs? 186 187 checkers = []188 if cond_concat_dims:189 c_in_channels = new_channels + cond_dim190 seq_length = in_height191 else:192 c_in_channels = new_channels193 seq_length = in_height + cond_seq_len194 for i in range(num_checkerboard):195 if norm_layer == "batchnorm":196 checkers += [BatchNorm(in_channels, bn_momentum)]197 elif norm_layer == "actnorm":198 checkers += [ActNorm(in_channels)]199 checkers += [InvertibleConv1x1(in_channels)]200 checkers += [Coupling(in_channels=c_in_channels,201 out_channels=in_channels,202 mid_channels=mid_channels,203 num_blocks=num_blocks,204 num_components=num_components,205 use_attn=use_attn,206 use_logmix=use_logmix,207 use_transformer_nn=use_transformer_nn,208 use_pos_emb=use_pos_emb,209 use_rel_pos_emb=use_rel_pos_emb,210 num_heads=num_heads,211 seq_length=seq_length,212 output_length=in_height,213 concat_dims=cond_concat_dims,214 drop_prob=drop_prob)]#,215 #Flip()]216 self.channels = nn.ModuleList(channels) if channels else None217 self.checkers = nn.ModuleList(checkers) if checkers else None218 219 if len(scales) <= 1:220 self.next = None221 else:222 next_shape = (in_channels, in_height // 2, in_width)223 self.next = _FlowStep(scales=scales[1:],224 in_shape=next_shape,225 cond_dim=2*cond_dim,226 mid_channels=mid_channels,227 num_blocks=num_blocks,228 num_components=num_components,229 use_attn=use_attn,230 use_logmix=use_logmix,231 use_transformer_nn=use_transformer_nn,232 use_pos_emb=use_pos_emb,233 use_rel_pos_emb=use_rel_pos_emb,234 num_heads=num_heads,235 norm_layer = norm_layer,236 bn_momentum = bn_momentum,237 cond_concat_dims = cond_concat_dims,238 cond_seq_len = cond_seq_len,239 drop_prob=drop_prob)240 241 self.z_shape = (in_channels, in_height, in_width)242 243 def z_dim(self):244 return self.z_shape245 246 def forward(self, x, cond, sldj, reverse=False):247 248 if reverse:249 #import pdb;pdb.set_trace()250 if self.next is not None:251 x = squeeze(x)252 cond = squeeze(cond)253 x, x_split = x.chunk(2, dim=1)254 x, sldj = self.next(x, cond, sldj, reverse)255 x = torch.cat((x, x_split), dim=1)256 x = unsqueeze(x)257 cond = unsqueeze(cond)258 259 if self.checkers:260 x = checkerboard(x)261 for flow in reversed(self.checkers):262 x, sldj = flow(x, cond, sldj, reverse)263 x = checkerboard(x, reverse=True)264 265 if self.channels:266 x = channelwise(x)267 for flow in reversed(self.channels):268 x, sldj = flow(x, cond, sldj, reverse)269 x = channelwise(x, reverse=True)270 else:271 # import pdb;pdb.set_trace()272 if self.channels:273 x = channelwise(x)274 for flow in self.channels:275 # import pdb;pdb.set_trace()276 x, sldj = flow(x, cond, sldj, reverse)277 # print(type(flow).__name__)278 # print(x[0].std())279 x = channelwise(x, reverse=True)280 281 if self.checkers:282 x = checkerboard(x)283 for flow in self.checkers:284 x, sldj = flow(x, cond, sldj, reverse)285 x = checkerboard(x, reverse=True)286 287 if self.next is not None:288 # import pdb;pdb.set_trace()289 # here we apply the flow steps but only to dimensions sampled at a lower scale. Hmm feels a bit weird290 x = squeeze(x)291 cond = squeeze(cond)292 x, x_split = x.chunk(2, dim=1)293 x, sldj = self.next(x, cond, sldj, reverse)294 x = torch.cat((x, x_split), dim=1)295 x = unsqueeze(x)296 297 # print(x.std())298 return x, sldj299 300