Spaces:
Build error
Build error
File size: 2,833 Bytes
f35cc94 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
from fastai.basics import *
from fastai.text.models.transformer import TransformerXL
from ..utils.attention_mask import rand_window_mask
class MusicTransformerXL(TransformerXL):
"Exactly like fastai's TransformerXL, but with more aggressive attention mask: see `rand_window_mask`"
def __init__(self, *args, encode_position=True, mask_steps=1, **kwargs):
import inspect
sig = inspect.signature(TransformerXL)
arg_params = { k:kwargs[k] for k in sig.parameters if k in kwargs }
super().__init__(*args, **arg_params)
self.encode_position = encode_position
if self.encode_position: self.beat_enc = BeatPositionEncoder(kwargs['d_model'])
self.mask_steps=mask_steps
def forward(self, x):
#The hidden state has to be initiliazed in the forward pass for nn.DataParallel
if self.mem_len > 0 and not self.init:
self.reset()
self.init = True
benc = 0
if self.encode_position:
x,pos = x['x'], x['pos']
benc = self.beat_enc(pos)
bs,x_len = x.size()
inp = self.drop_emb(self.encoder(x) + benc) #.mul_(self.d_model ** 0.5)
m_len = self.hidden[0].size(1) if hasattr(self, 'hidden') and len(self.hidden[0].size()) > 1 else 0
seq_len = m_len + x_len
mask = rand_window_mask(x_len, m_len, inp.device, max_size=self.mask_steps, is_eval=not self.training) if self.mask else None
if m_len == 0: mask[...,0,0] = 0
#[None,:,:None] for einsum implementation of attention
hids = []
pos = torch.arange(seq_len-1, -1, -1, device=inp.device, dtype=inp.dtype)
pos_enc = self.pos_enc(pos)
hids.append(inp)
for i, layer in enumerate(self.layers):
mem = self.hidden[i] if self.mem_len > 0 else None
inp = layer(inp, r=pos_enc, u=self.u, v=self.v, mask=mask, mem=mem)
hids.append(inp)
core_out = inp[:,-x_len:]
if self.mem_len > 0 : self._update_mems(hids)
return (self.hidden if self.mem_len > 0 else [core_out]),[core_out]
# Beat encoder
class BeatPositionEncoder(nn.Module):
"Embedding + positional encoding + dropout"
def __init__(self, emb_sz:int, beat_len=32, max_bar_len=1024):
super().__init__()
self.beat_len, self.max_bar_len = beat_len, max_bar_len
self.beat_enc = nn.Embedding(beat_len, emb_sz, padding_idx=0)
self.bar_enc = nn.Embedding(max_bar_len, emb_sz, padding_idx=0)
def forward(self, pos):
beat_enc = self.beat_enc(pos % self.beat_len)
bar_pos = pos // self.beat_len % self.max_bar_len
bar_pos[bar_pos >= self.max_bar_len] = self.max_bar_len - 1
bar_enc = self.bar_enc((bar_pos))
return beat_enc + bar_enc |