refactor

.gitignore

@@ -171,3 +171,5 @@ archived/
 scratch/
 
 runs-archive
+lyrebird-audiotools
+lyrebird-audio-codec

vampnet/modules/activations.py

@@ -0,0 +1,55 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+
+
+class NewGELU(nn.Module):
+    """
+    Implementation of the GELU activation function currently in Google BERT repo
+    (identical to OpenAI GPT). Also see the Gaussian Error Linear Units
+    paper: https://arxiv.org/abs/1606.08415
+    """
+
+    def forward(self, x):
+        return (
+            0.5
+            * x
+            * (
+                1.0
+                + torch.tanh(
+                    math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))
+                )
+            )
+        )
+
+class GatedGELU(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gelu = NewGELU()
+
+    def forward(self, x, dim: int = -1):
+        p1, p2 = x.chunk(2, dim=dim)
+        return p1 * self.gelu(p2)
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(channels))
+
+    def forward(self, x):
+        return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
+
+def get_activation(name: str = "relu"):
+    if name == "relu":
+        return nn.ReLU
+    elif name == "gelu":
+        return NewGELU
+    elif name == "geglu":
+        return GatedGELU
+    elif name == "snake":
+        return Snake1d
+    else:
+        raise ValueError(f"Unrecognized activation {name}")

vampnet/modules/{modules.py → layers.py}

@@ -26,25 +26,6 @@ def recurse_children(module, fn):
         yield fn(child)
 
 
-# Scripting this brings model speed up 1.4x
-@torch.jit.script
-def snake(x, alpha):
-    shape = x.shape
-    x = x.reshape(shape[0], shape[1], -1)
-    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
-    x = x.reshape(shape)
-    return x
-
-
-class Snake1d(nn.Module):
-    def __init__(self, channels):
-        super().__init__()
-        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
-
-    def forward(self, x):
-        return snake(x, self.alpha)
-
-
 def WNConv1d(*args, **kwargs):
     return weight_norm(nn.Conv1d(*args, **kwargs))
 

vampnet/modules/transformer.py

@@ -7,10 +7,11 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from .base import VampBase
-from .modules import CodebookEmbedding
-from .modules import FiLM
-from .modules import SequentialWithFiLM
-from .modules import WNConv1d
+from .activations import get_activation
+from .layers import CodebookEmbedding
+from .layers import FiLM
+from .layers import SequentialWithFiLM
+from .layers import WNConv1d
 
 
 class RMSNorm(nn.Module):
@@ -37,58 +38,6 @@ class RMSNorm(nn.Module):
         return self.weight * x
 
 
-def get_activation(name: str = "relu"):
-    if name == "relu":
-        return nn.ReLU
-    elif name == "gelu":
-        return NewGELU
-    elif name == "geglu":
-        return GatedGELU
-    elif name == "snake":
-        return Snake1d
-    else:
-        raise ValueError(f"Unrecognized activation {name}")
-
-
-class NewGELU(nn.Module):
-    """
-    Implementation of the GELU activation function currently in Google BERT repo
-    (identical to OpenAI GPT). Also see the Gaussian Error Linear Units
-    paper: https://arxiv.org/abs/1606.08415
-    """
-
-    def forward(self, x):
-        return (
-            0.5
-            * x
-            * (
-                1.0
-                + torch.tanh(
-                    math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))
-                )
-            )
-        )
-
-
-class GatedGELU(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.gelu = NewGELU()
-
-    def forward(self, x, dim: int = -1):
-        p1, p2 = x.chunk(2, dim=dim)
-        return p1 * self.gelu(p2)
-
-
-class Snake1d(nn.Module):
-    def __init__(self, channels):
-        super().__init__()
-        self.alpha = nn.Parameter(torch.ones(channels))
-
-    def forward(self, x):
-        return x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
-
-
 class FeedForward(nn.Module):
     def __init__(
         self, d_model: int = 512, dropout: float = 0.1, activation: str = "geglu"
@@ -572,7 +521,7 @@ class VampNet(VampBase):
 
 if __name__ == "__main__":
     # import argbind
-    from .modules import num_params
+    from .layers import num_params
 
     VampNet = argbind.bind(VampNet)