make sure code expiry feature works with cosine sim

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -104,9 +104,10 @@ def expire_codes_(self, batch_samples):
             return
 
         expired_codes = self.cluster_size < self.threshold_ema_dead_code
-        if torch.any(expired_codes):
-            batch_samples = rearrange(batch_samples, '... d -> (...) d')
-            self.replace(batch_samples, mask = expired_codes)
+        if not torch.any(expired_codes):
+            return
+        batch_samples = rearrange(batch_samples, '... d -> (...) d')
+        self.replace(batch_samples, mask = expired_codes)
 
     def forward(self, x):
         shape, dtype = x.shape, x.dtype
@@ -163,6 +164,7 @@ def __init__(
         self.threshold_ema_dead_code = threshold_ema_dead_code
 
         self.register_buffer('initted', torch.Tensor([not kmeans_init]))
+        self.register_buffer('cluster_size', torch.zeros(codebook_size))
         self.register_buffer('embed', embed)
 
     def init_embed_(self, data):
@@ -185,9 +187,10 @@ def expire_codes_(self, batch_samples):
             return
 
         expired_codes = self.cluster_size < self.threshold_ema_dead_code
-        if torch.any(expired_codes):
-            batch_samples = rearrange(batch_samples, '... d -> (...) d')
-            self.replace(batch_samples, mask = expired_codes)
+        if not torch.any(expired_codes):
+            return
+        batch_samples = rearrange(batch_samples, '... d -> (...) d')
+        self.replace(batch_samples, mask = expired_codes)
 
     def forward(self, x):
         shape, dtype = x.shape, x.dtype
@@ -207,6 +210,8 @@ def forward(self, x):
 
         if self.training:
             bins = embed_onehot.sum(0)
+            ema_inplace(self.cluster_size, bins, self.decay)
+
             zero_mask = (bins == 0)
             bins = bins.masked_fill(zero_mask, 1.)