add velocity and multi-factorization

Author: victor-shepardson

notepredictor/notepredictor/data.py

@@ -28,6 +28,7 @@ def __getitem__(self, idx):
         item = torch.load(f)
         pitch = item['pitch'] # 1-d LongTensor of MIDI pitches 0-127
         time = item['time']
+        velocity = item['velocity']
         assert len(pitch) == len(time)
 
         # random transpose avoiding out of range notes
@@ -45,6 +46,11 @@ def __getitem__(self, idx):
             time + (torch.rand_like(time)-0.5)*2e-3
             ).clamp(0., float('inf'))
 
+        velocity = (
+            velocity + 
+            (torch.rand_like(time)-0.5) * ((velocity>0) & (velocity<127)).float()
+            ).clamp(0., 127.)
+
         # pad with start, end tokens
         pad = max(1, self.batch_len-len(pitch))
         pitch = torch.cat((
@@ -55,15 +61,21 @@ def __getitem__(self, idx):
             time.new_zeros((1,)),
             time,
             time.new_zeros((pad,))))
+        velocity = torch.cat((
+            velocity.new_zeros((1,)),
+            velocity,
+            velocity.new_zeros((pad,))))
 
         # random slice
         i = random.randint(0, len(pitch)-self.batch_len)
         pitch = pitch[i:i+self.batch_len]
         time = time[i:i+self.batch_len]
+        velocity = velocity[i:i+self.batch_len]
 
         # time = time.clamp(*self.clamp_time)
 
         return {
             'pitch':pitch,
-            'time':time
+            'time':time,
+            'velocity':velocity
         }

notepredictor/notepredictor/distributions.py

@@ -6,20 +6,28 @@
 import torch.nn.functional as F
 
 class CensoredMixtureLogistic(nn.Module):
-    def __init__(self, n, res=1e-2, lo='-inf', hi='inf', sharp_bounds=(1e-4,2e3)):
+    def __init__(self, n, res=1e-2, lo='-inf', hi='inf', 
+            sharp_bounds=(1e-4,2e3), init=None):
         super().__init__()
         self.n = n
         self.res = res
         self.sharp_bounds = sharp_bounds
-        # self.base_bounds = tuple(reversed([math.exp(-s) for s in sharp_bounds]))
-        # self.register_buffer('max_sharp', torch.tensor(float(max_sharp)))
         self.register_buffer('lo', torch.tensor(float(lo)))
         self.register_buffer('hi', torch.tensor(float(hi)))
         # TODO: init is not general-purpose
         #TODO
-        self.bias = nn.Parameter(torch.cat((
-            torch.zeros(n), torch.logspace(-3,1,n), torch.zeros(n)
-            )))
+        if init=='time':
+            self.bias = nn.Parameter(torch.cat((
+                torch.zeros(n), torch.logspace(-3,1,n), torch.zeros(n)
+                )))
+        elif init=='velocity':
+            self.bias = nn.Parameter(torch.cat((
+                torch.zeros(n), torch.linspace(1,126,n), torch.zeros(n)
+                )))
+        else:
+            self.bias = nn.Parameter(torch.cat((
+                torch.zeros(n), torch.randn(n), torch.zeros(n)
+                )))
 
     @property
     def n_params(self):
@@ -28,7 +36,7 @@ def n_params(self):
     def get_params(self, h):
         assert h.shape[-1] == self.n_params
         h = h+self.bias
-        # get parameters fron unconstrained hidden state:
+        # get parameters from unconstrained hidden state:
         logit_pi, loc, log_s = torch.chunk(h, 3, -1)
         # mixture coefficients
         log_pi = logit_pi - logit_pi.logsumexp(-1,keepdim=True)
@@ -37,48 +45,30 @@ def get_params(self, h):
         # sharpness
         s = F.softplus(log_s).clamp(*self.sharp_bounds)
         return log_pi, loc, s
-        # exp of negative sharpness
-        # b = logit_b.sigmoid().clamp(*self.base_bounds)
-        # log_b = F.logsigmoid(logit_b).clamp(
-            # math.log(self.base_bounds[0]), math.log(self.base_bounds[1]))
-        # b = ((logit_b / (logit_b.abs() + 1))*.5+.5).clamp(*self.base_bounds)
-        # log_b = b.log()
-        # return log_pi, loc, b, log_b
+
 
     def forward(self, h, x):
         """log prob of x under distribution parameterized by h"""
-        # log_pi, loc, b, log_b = self.get_params(h)
         log_pi, loc, s = self.get_params(h)    
 
         d = self.res/2
-        x = x.clamp(self.lo-d, self.hi+d)[...,None]
+        x = x.clamp(self.lo, self.hi)[...,None]
         x_ = (x - loc) * s
-
-        # numerical crimes follow
-        # q = b ** -x_
-        q = x_.exp()
         sd = s*d
-        # bdp, bdm = b**d, b**-d
-        sdm, sdp = (-sd).exp(), sd.exp()
-        # # censoring
-        lo_cens = x <= self.lo
-        hi_cens = x >= self.hi
-        ones = torch.ones_like(q)
-        zeros = torch.zeros_like(q)
-
-        diff_term = torch.where(
-            lo_cens | hi_cens, ones, sdp - sdm).log()
-        # sdm_term = torch.where(
-        #     hi_cens, ones, (q + sdm)).log()
-        sdm_term = torch.where(hi_cens, zeros, x_ + F.softplus(-sd-x_))
-        # sdp_term = torch.where(
-        #     lo_cens, ones, (q + sdp)).log()
-        sdp_term = torch.where(lo_cens, zeros, x_ + F.softplus(sd-x_))
-        x_or_sd = torch.where(hi_cens, sd, x_)
 
-        log_delta_cdf = (
-            x_or_sd + diff_term - sdm_term - sdp_term
-        )
+        # # censoring
+        lo_cens = x <= self.lo+d
+        hi_cens = x >= self.hi-d
+        ones = torch.ones_like(x_)
+        zeros = torch.zeros_like(x_)
+
+        diff_term = torch.where(lo_cens | hi_cens, 
+            ones, sd.exp() - (-sd).exp()
+            ).log()
+        minus_sp_term = torch.where(hi_cens, -sd, F.softplus(-sd-x_))
+        plus_sp_term = torch.where(lo_cens, zeros, x_ + F.softplus(sd-x_))
+
+        log_delta_cdf = diff_term - minus_sp_term - plus_sp_term
 
         # log prob
         r = {
@@ -87,23 +77,23 @@ def forward(self, h, x):
         # diagnostics
         with torch.no_grad():
             ent = D.Categorical(logits=log_pi).entropy()
-            # s = -b.log()
             r |= {
-                'min_sharpness': s.min(),
+                # 'min_sharpness': s.min(),
                 'max_sharpness': s.max(),
-                'min_entropy': ent.min(),
-                'max_entropy': ent.max(),
-                'marginal_entropy': D.Categorical(
+                'mean_sharpness': (s*log_pi.exp()).sum(-1).mean(),
+                # 'min_entropy': ent.min(),
+                # 'max_entropy': ent.max(),
+                'mean_cmp_entropy': ent.mean(),
+                'marginal_cmp_entropy': D.Categorical(
                     log_pi.exp().mean(list(range(log_pi.ndim-1)))).entropy(),
-                'min_loc': loc.min(),
-                'max_loc': loc.max()
+                # 'min_loc': loc.min(),
+                # 'max_loc': loc.max()
             }
         return r
 
     def cdf(self, h, x):
         log_pi, loc, s = self.get_params(h)  
         x_ = (x[...,None] - loc) * s
-        # cdfs = 1 / (1 + b ** x_)
         cdfs = x_.sigmoid()
         cdf = (cdfs * log_pi.softmax(-1)).sum(-1)
         return cdf
@@ -113,9 +103,7 @@ def sample(self, h, shape=1):
         Args:
             shape: additional sample shape to be prepended to dims
         """
-        # log_pi, loc, _, log_b = self.get_params(h)
         log_pi, loc, s = self.get_params(h)
-        # scale = -1/log_b
         scale = 1/s
 
         c = D.Categorical(logits=log_pi).sample((shape,))