uniform init the codebooks, thanks to @ramyamounir for pointing this out

lucidrains · Jul 11, 2022 · abbc41a · abbc41a
1 parent dea416f
commit abbc41a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '0.8.0',
+  version = '0.8.1',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   author = 'Phil Wang',

diff --git a/vector_quantize_pytorch/vector_quantize_pytorch.py b/vector_quantize_pytorch/vector_quantize_pytorch.py
@@ -22,6 +22,11 @@ def l2norm(t):
 def log(t, eps = 1e-20):
     return torch.log(t.clamp(min = eps))
 
+def uniform_init(*shape):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
 def gumbel_noise(t):
     noise = torch.zeros_like(t).uniform_(0, 1)
     return -log(-log(noise))
@@ -162,7 +167,7 @@ def __init__(
     ):
         super().__init__()
         self.decay = decay
-        init_fn = torch.randn if not kmeans_init else torch.zeros
+        init_fn = uniform_init if not kmeans_init else torch.zeros
         embed = init_fn(codebook_size, dim)
 
         self.codebook_size = codebook_size
@@ -268,7 +273,7 @@ def __init__(
         self.decay = decay
 
         if not kmeans_init:
-            embed = l2norm(torch.randn(codebook_size, dim))
+            embed = l2norm(uniform_init(codebook_size, dim))
         else:
             embed = torch.zeros(codebook_size, dim)