Spaces:
Running
Running
| import torch | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from transformer_encoder import VAETransformerEncoder | |
| from transformer_helpers import ( | |
| weights_init, PositionalEncoding, TokenEmbedding, generate_causal_mask | |
| ) | |
| class VAETransformerDecoder(nn.Module): | |
| def __init__(self, n_layer, n_head, d_model, d_ff, d_seg_emb, dropout=0.1, activation='relu', cond_mode='in-attn'): | |
| super(VAETransformerDecoder, self).__init__() | |
| self.n_layer = n_layer | |
| self.n_head = n_head | |
| self.d_model = d_model | |
| self.d_ff = d_ff | |
| self.d_seg_emb = d_seg_emb | |
| self.dropout = dropout | |
| self.activation = activation | |
| self.cond_mode = cond_mode | |
| if cond_mode == 'in-attn': | |
| self.seg_emb_proj = nn.Linear(d_seg_emb, d_model, bias=False) | |
| elif cond_mode == 'pre-attn': | |
| self.seg_emb_proj = nn.Linear(d_seg_emb + d_model, d_model, bias=False) | |
| self.decoder_layers = nn.ModuleList() | |
| for i in range(n_layer): | |
| self.decoder_layers.append( | |
| nn.TransformerEncoderLayer(d_model, n_head, d_ff, dropout, activation) | |
| ) | |
| def forward(self, x, seg_emb): | |
| if not hasattr(self, 'cond_mode'): | |
| self.cond_mode = 'in-attn' | |
| attn_mask = generate_causal_mask(x.size(0)).to(x.device) | |
| # print (attn_mask.size()) | |
| if self.cond_mode == 'in-attn': | |
| seg_emb = self.seg_emb_proj(seg_emb) | |
| elif self.cond_mode == 'pre-attn': | |
| x = torch.cat([x, seg_emb], dim=-1) | |
| x = self.seg_emb_proj(x) | |
| out = x | |
| for i in range(self.n_layer): | |
| if self.cond_mode == 'in-attn': | |
| out += seg_emb | |
| out = self.decoder_layers[i](out, src_mask=attn_mask) | |
| return out | |
| class MuseMorphose(nn.Module): | |
| def __init__(self, enc_n_layer, enc_n_head, enc_d_model, enc_d_ff, | |
| dec_n_layer, dec_n_head, dec_d_model, dec_d_ff, | |
| d_vae_latent, d_embed, n_token, | |
| enc_dropout=0.1, enc_activation='relu', | |
| dec_dropout=0.1, dec_activation='relu', | |
| d_rfreq_emb=32, d_polyph_emb=32, | |
| n_rfreq_cls=8, n_polyph_cls=8, | |
| is_training=True, use_attr_cls=True, | |
| cond_mode='in-attn' | |
| ): | |
| super(MuseMorphose, self).__init__() | |
| self.enc_n_layer = enc_n_layer | |
| self.enc_n_head = enc_n_head | |
| self.enc_d_model = enc_d_model | |
| self.enc_d_ff = enc_d_ff | |
| self.enc_dropout = enc_dropout | |
| self.enc_activation = enc_activation | |
| self.dec_n_layer = dec_n_layer | |
| self.dec_n_head = dec_n_head | |
| self.dec_d_model = dec_d_model | |
| self.dec_d_ff = dec_d_ff | |
| self.dec_dropout = dec_dropout | |
| self.dec_activation = dec_activation | |
| self.d_vae_latent = d_vae_latent | |
| self.n_token = n_token | |
| self.is_training = is_training | |
| self.cond_mode = cond_mode | |
| self.token_emb = TokenEmbedding(n_token, d_embed, enc_d_model) | |
| self.d_embed = d_embed | |
| self.pe = PositionalEncoding(d_embed) | |
| self.dec_out_proj = nn.Linear(dec_d_model, n_token) | |
| self.encoder = VAETransformerEncoder( | |
| enc_n_layer, enc_n_head, enc_d_model, enc_d_ff, d_vae_latent, enc_dropout, enc_activation | |
| ) | |
| self.use_attr_cls = use_attr_cls | |
| if use_attr_cls: | |
| self.decoder = VAETransformerDecoder( | |
| dec_n_layer, dec_n_head, dec_d_model, dec_d_ff, d_vae_latent + d_polyph_emb + d_rfreq_emb, | |
| dropout=dec_dropout, activation=dec_activation, | |
| cond_mode=cond_mode | |
| ) | |
| else: | |
| self.decoder = VAETransformerDecoder( | |
| dec_n_layer, dec_n_head, dec_d_model, dec_d_ff, d_vae_latent, | |
| dropout=dec_dropout, activation=dec_activation, | |
| cond_mode=cond_mode | |
| ) | |
| if use_attr_cls: | |
| self.d_rfreq_emb = d_rfreq_emb | |
| self.d_polyph_emb = d_polyph_emb | |
| self.rfreq_attr_emb = TokenEmbedding(n_rfreq_cls, d_rfreq_emb, d_rfreq_emb) | |
| self.polyph_attr_emb = TokenEmbedding(n_polyph_cls, d_polyph_emb, d_polyph_emb) | |
| else: | |
| self.rfreq_attr_emb = None | |
| self.polyph_attr_emb = None | |
| self.emb_dropout = nn.Dropout(self.enc_dropout) | |
| self.apply(weights_init) | |
| def reparameterize(self, mu, logvar, use_sampling=True, sampling_var=1.): | |
| std = torch.exp(0.5 * logvar).to(mu.device) | |
| if use_sampling: | |
| eps = torch.randn_like(std).to(mu.device) * sampling_var | |
| else: | |
| eps = torch.zeros_like(std).to(mu.device) | |
| return eps * std + mu | |
| def get_sampled_latent(self, inp, padding_mask=None, use_sampling=False, sampling_var=0.): | |
| token_emb = self.token_emb(inp) | |
| enc_inp = self.emb_dropout(token_emb) + self.pe(inp.size(0)) | |
| _, mu, logvar = self.encoder(enc_inp, padding_mask=padding_mask) | |
| mu, logvar = mu.reshape(-1, mu.size(-1)), logvar.reshape(-1, mu.size(-1)) | |
| vae_latent = self.reparameterize(mu, logvar, use_sampling=use_sampling, sampling_var=sampling_var) | |
| return vae_latent | |
| def generate(self, inp, dec_seg_emb, rfreq_cls=None, polyph_cls=None, keep_last_only=True): | |
| token_emb = self.token_emb(inp) | |
| dec_inp = self.emb_dropout(token_emb) + self.pe(inp.size(0)) | |
| if rfreq_cls is not None and polyph_cls is not None: | |
| dec_rfreq_emb = self.rfreq_attr_emb(rfreq_cls) | |
| dec_polyph_emb = self.polyph_attr_emb(polyph_cls) | |
| dec_seg_emb_cat = torch.cat([dec_seg_emb, dec_rfreq_emb, dec_polyph_emb], dim=-1) | |
| else: | |
| dec_seg_emb_cat = dec_seg_emb | |
| out = self.decoder(dec_inp, dec_seg_emb_cat) | |
| out = self.dec_out_proj(out) | |
| if keep_last_only: | |
| out = out[-1, ...] | |
| return out | |
| def forward(self, enc_inp, dec_inp, dec_inp_bar_pos, rfreq_cls=None, polyph_cls=None, padding_mask=None): | |
| # [shape of enc_inp] (seqlen_per_bar, bsize, n_bars_per_sample) | |
| enc_bt_size, enc_n_bars = enc_inp.size(1), enc_inp.size(2) | |
| enc_token_emb = self.token_emb(enc_inp) | |
| # [shape of dec_inp] (seqlen_per_sample, bsize) | |
| # [shape of rfreq_cls & polyph_cls] same as above | |
| # -- (should copy each bar's label to all corresponding indices) | |
| dec_token_emb = self.token_emb(dec_inp) | |
| enc_token_emb = enc_token_emb.reshape( | |
| enc_inp.size(0), -1, enc_token_emb.size(-1) | |
| ) | |
| enc_inp = self.emb_dropout(enc_token_emb) + self.pe(enc_inp.size(0)) | |
| dec_inp = self.emb_dropout(dec_token_emb) + self.pe(dec_inp.size(0)) | |
| # [shape of padding_mask] (bsize, n_bars_per_sample, seqlen_per_bar) | |
| # -- should be `True` for padded indices (i.e., those >= seqlen of the bar), `False` otherwise | |
| if padding_mask is not None: | |
| padding_mask = padding_mask.reshape(-1, padding_mask.size(-1)) | |
| _, mu, logvar = self.encoder(enc_inp, padding_mask=padding_mask) | |
| vae_latent = self.reparameterize(mu, logvar) | |
| vae_latent_reshaped = vae_latent.reshape(enc_bt_size, enc_n_bars, -1) | |
| dec_seg_emb = torch.zeros(dec_inp.size(0), dec_inp.size(1), self.d_vae_latent).to(vae_latent.device) | |
| for n in range(dec_inp.size(1)): | |
| # [shape of dec_inp_bar_pos] (bsize, n_bars_per_sample + 1) | |
| # -- stores [[start idx of bar #1, sample #1, ..., start idx of bar #K, sample #1, seqlen of sample #1], [same for another sample], ...] | |
| for b, (st, ed) in enumerate(zip(dec_inp_bar_pos[n, :-1], dec_inp_bar_pos[n, 1:])): | |
| dec_seg_emb[st:ed, n, :] = vae_latent_reshaped[n, b, :] | |
| if rfreq_cls is not None and polyph_cls is not None and self.use_attr_cls: | |
| dec_rfreq_emb = self.rfreq_attr_emb(rfreq_cls) | |
| dec_polyph_emb = self.polyph_attr_emb(polyph_cls) | |
| dec_seg_emb_cat = torch.cat([dec_seg_emb, dec_rfreq_emb, dec_polyph_emb], dim=-1) | |
| else: | |
| dec_seg_emb_cat = dec_seg_emb | |
| dec_out = self.decoder(dec_inp, dec_seg_emb_cat) | |
| dec_logits = self.dec_out_proj(dec_out) | |
| return mu, logvar, dec_logits | |
| def compute_loss(self, mu, logvar, beta, fb_lambda, dec_logits, dec_tgt): | |
| recons_loss = F.cross_entropy( | |
| dec_logits.view(-1, dec_logits.size(-1)), dec_tgt.contiguous().view(-1), | |
| ignore_index=self.n_token - 1, reduction='mean' | |
| ).float() | |
| kl_raw = -0.5 * (1 + logvar - mu ** 2 - logvar.exp()).mean(dim=0) | |
| kl_before_free_bits = kl_raw.mean() | |
| kl_after_free_bits = kl_raw.clamp(min=fb_lambda) | |
| kldiv_loss = kl_after_free_bits.mean() | |
| return { | |
| 'beta': beta, | |
| 'total_loss': recons_loss + beta * kldiv_loss, | |
| 'kldiv_loss': kldiv_loss, | |
| 'kldiv_raw': kl_before_free_bits, | |
| 'recons_loss': recons_loss | |
| } |