make init param more intuitive, but leave n_embed backwards compatible

Author: lucidrains

README.md

@@ -18,9 +18,9 @@ from vector_quantize_pytorch import VectorQuantize
 
 vq = VectorQuantize(
     dim = 256,
-    n_embed = 512,     # size of the dictionary
-    decay = 0.8,       # the exponential moving average decay, lower means the dictionary will change faster
-    commitment = 1.    # the weight on the commitment loss
+    codebook_size = 512,     # codebook size
+    decay = 0.8,             # the exponential moving average decay, lower means the dictionary will change faster
+    commitment = 1.          # the weight on the commitment loss
 )
 
 x = torch.randn(1, 1024, 256)
@@ -38,7 +38,7 @@ from vector_quantize_pytorch import ResidualVQ
 residual_vq = ResidualVQ(
     dim = 256,
     num_quantizers = 8,      # specify number of quantizers
-    n_embed = 1024,          # codebook size
+    codebook_size = 1024,    # codebook size
 )
 
 x = torch.randn(1, 1024, 256)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '0.2.0',
+  version = '0.2.1',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   author = 'Phil Wang',

vector_quantize_pytorch/residual_vq.py

@@ -8,11 +8,10 @@ def __init__(
         self,
         *,
         num_quantizers,
-        n_embed,
         **kwargs
     ):
         super().__init__()
-        self.layers = nn.ModuleList([VectorQuantize(n_embed = n_embed, **kwargs) for _ in range(num_quantizers)])
+        self.layers = nn.ModuleList([VectorQuantize(**kwargs) for _ in range(num_quantizers)])
 
     def forward(self, x):
         quantized_out = 0.

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -2,6 +2,12 @@
 from torch import nn
 import torch.nn.functional as F
 
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
 def ema_inplace(moving_avg, new, decay):
     moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))
 
@@ -12,12 +18,14 @@ class VectorQuantize(nn.Module):
     def __init__(
         self,
         dim,
-        n_embed,
+        codebook_size,
         decay = 0.8,
         commitment = 1.,
-        eps = 1e-5
+        eps = 1e-5,
+        n_embed = None,
     ):
         super().__init__()
+        n_embed = default(n_embed, codebook_size)
 
         self.dim = dim
         self.n_embed = n_embed