Dead code replacement is actually based on the exponential moving average

Author: wesbz

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -67,7 +67,8 @@ def __init__(
         kmeans_init = False,
         kmeans_iters = 10,
         decay = 0.8,
-        eps = 1e-5
+        eps = 1e-5,
+        threshold_ema_dead_code = 2
     ):
         super().__init__()
         self.decay = decay
@@ -77,6 +78,7 @@ def __init__(
         self.codebook_size = codebook_size
         self.kmeans_iters = kmeans_iters
         self.eps = eps
+        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))
@@ -97,6 +99,15 @@ def replace(self, samples, mask):
         )
         self.embed.data.copy_(modified_codebook)
 
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            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)
+
     def forward(self, x):
         shape, dtype = x.shape, x.dtype
         flatten = rearrange(x, '... d -> (...) d')
@@ -112,7 +123,7 @@ def forward(self, x):
         )
 
         embed_ind = dist.max(dim = -1).indices
-        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(x.dtype)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
         embed_ind = embed_ind.view(*shape[:-1])
         quantize = F.embedding(embed_ind, self.embed)
 
@@ -123,6 +134,7 @@ def forward(self, x):
             cluster_size = laplace_smoothing(self.cluster_size, self.codebook_size, self.eps) * self.cluster_size.sum()
             embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
             self.embed.data.copy_(embed_normalized)
+            self.expire_codes_(x)
 
         return quantize, embed_ind
 
@@ -134,7 +146,8 @@ def __init__(
         kmeans_init = False,
         kmeans_iters = 10,
         decay = 0.8,
-        eps = 1e-5
+        eps = 1e-5,
+        threshold_ema_dead_code = 2
     ):
         super().__init__()
         self.decay = decay
@@ -147,6 +160,7 @@ def __init__(
         self.codebook_size = codebook_size
         self.kmeans_iters = kmeans_iters
         self.eps = eps
+        self.threshold_ema_dead_code = threshold_ema_dead_code
 
         self.register_buffer('initted', torch.Tensor([not kmeans_init]))
         self.register_buffer('embed', embed)
@@ -166,6 +180,15 @@ def replace(self, samples, mask):
         )
         self.embed.data.copy_(modified_codebook)
 
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            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)
+
     def forward(self, x):
         shape, dtype = x.shape, x.dtype
         flatten = rearrange(x, '... d -> (...) d')
@@ -193,6 +216,7 @@ def forward(self, x):
             embed_normalized = torch.where(zero_mask[..., None], embed,
                                            embed_normalized)
             ema_inplace(self.embed, embed_normalized, self.decay)
+            self.expire_codes_(x)
 
         return quantize, embed_ind
 
@@ -211,7 +235,7 @@ def __init__(
         kmeans_init = False,
         kmeans_iters = 10,
         use_cosine_sim = False,
-        max_codebook_misses_before_expiry = 0
+        threshold_ema_dead_code = 0
     ):
         super().__init__()
         n_embed = default(n_embed, codebook_size)
@@ -229,44 +253,23 @@ def __init__(
         codebook_class = EuclideanCodebook if not use_cosine_sim \
                          else CosineSimCodebook
 
-        self._codebook = klass(
+        self._codebook = codebook_class(
             dim = codebook_dim,
             codebook_size = n_embed,
             kmeans_init = kmeans_init,
             kmeans_iters = kmeans_iters,
             decay = decay,
-            eps = eps
+            eps = eps,
+            threshold_ema_dead_code = threshold_ema_dead_code
         )
 
         self.codebook_size = codebook_size
-        self.max_codebook_misses_before_expiry = max_codebook_misses_before_expiry
-
-        if max_codebook_misses_before_expiry > 0:
-            codebook_misses = torch.zeros(codebook_size)
-            self.register_buffer('codebook_misses', codebook_misses)
 
     @property
     def codebook(self):
         return self._codebook.codebook
 
-    def expire_codes_(self, embed_ind, batch_samples):
-        if self.max_codebook_misses_before_expiry == 0:
-            return
-
-        embed_ind = rearrange(embed_ind, '... -> (...)')
-        misses = torch.bincount(embed_ind, minlength = self.codebook_size) == 0
-        self.codebook_misses += misses
-
-        expired_codes = self.codebook_misses >= self.max_codebook_misses_before_expiry
-        if not torch.any(expired_codes):
-            return
-
-        self.codebook_misses.masked_fill_(expired_codes, 0)
-        batch_samples = rearrange(batch_samples, '... d -> (...) d')
-        self._codebook.replace(batch_samples, mask = expired_codes)
-
     def forward(self, x):
-        dtype = x.dtype
         x = self.project_in(x)
 
         quantize, embed_ind = self._codebook(x)
@@ -276,7 +279,6 @@ def forward(self, x):
         if self.training:
             commit_loss = F.mse_loss(quantize.detach(), x) * self.commitment
             quantize = x + (quantize - x).detach()
-            self.expire_codes_(embed_ind, x)
 
         quantize = self.project_out(quantize)
         return quantize, embed_ind, commit_loss