vector-quantize-pytorch 1.22.18__tar.gz → 1.23.4__tar.gz

{vector_quantize_pytorch-1.22.18 → vector_quantize_pytorch-1.23.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.22.18
+Version: 1.23.4
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch

{vector_quantize_pytorch-1.22.18 → vector_quantize_pytorch-1.23.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.22.18"
+version = "1.23.4"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{vector_quantize_pytorch-1.22.18 → vector_quantize_pytorch-1.23.4}/tests/test_readme.py

@@ -434,3 +434,30 @@ def test_vq_custom_ema_update_weighting(
 
     assert torch.allclose(codebook_before[did_not_update], codebook_after[did_not_update], atol = 1e-6)
     assert (codebook_before[did_update] != codebook_after[did_update]).all()
+
+def test_accum_ema_update():
+    from vector_quantize_pytorch import VectorQuantize
+
+    vq = VectorQuantize(
+        dim = 256,
+        use_cosine_sim = True,
+        codebook_dim = 128,
+        codebook_size = 8,        # codebook size
+        decay = 0.8,              # the exponential moving average decay, lower means the dictionary will change faster
+        commitment_weight = 1.,   # the weight on the commitment loss
+    )
+
+    x = torch.randn(16, 1024, 256)
+
+    codebook_before = vq.codebook.clone()
+
+    vq.train()
+
+    _ = vq(x, accum_ema_update = True)
+    _ = vq(x, accum_ema_update = True)
+
+    assert torch.allclose(codebook_before, vq.codebook, atol = 1e-6)
+
+    _ = vq(x)
+
+    assert not torch.allclose(codebook_before, vq.codebook, atol = 1e-6)

{vector_quantize_pytorch-1.22.18 → vector_quantize_pytorch-1.23.4}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -61,7 +61,7 @@ def floor_ste(z):
 class FSQ(Module):
     def __init__(
         self,
-        levels: list[int],
+        levels: list[int] | tuple[int, ...],
         dim: int | None = None,
         num_codebooks = 1,
         keep_num_codebooks_dim: bool | None = None,
@@ -76,6 +76,9 @@ class FSQ(Module):
     ):
         super().__init__()
 
+        if isinstance(levels, tuple):
+            levels = list(levels)
+
         _levels = tensor(levels, dtype = int32)
         self.register_buffer('_levels', _levels, persistent = False)
 

{vector_quantize_pytorch-1.22.18 → vector_quantize_pytorch-1.23.4}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -61,7 +61,22 @@ def log(t, eps = 1e-20):
 def entropy(prob, eps = 1e-5):
     return (-prob * log(prob, eps = eps)).sum(dim = -1)
 
+def accum_grad_(t, grad):
+    if exists(t.grad):
+        t.grad.add_(grad)
+    else:
+        t.grad = grad.clone().detach()
+
 def ema_inplace(old, new, decay, weight = None):
+
+    # if old.grad is populated, add it to new and set it to None
+
+    if exists(old.grad):
+        new.add_(old.grad)
+        old.grad = None
+
+    # take care of custom weighting
+
     weight = default(weight, 1.)
 
     if is_tensor(weight):
@@ -71,7 +86,7 @@ def ema_inplace(old, new, decay, weight = None):
         assert weight.ndim == 2 and weight.shape == old.shape[:2]
         weight = append_dims_to(weight, old.ndim)
 
-    old.lerp_(new, (1. - decay) * weight)
+    old.data.lerp_(new, (1. - decay) * weight)
 
 def pack_one(t, pattern):
     packed, ps = pack([t], pattern)
@@ -258,12 +273,14 @@ def efficient_rotation_trick_transform(u, q, e):
     e = rearrange(e, 'b d -> b 1 d')
     w = l2norm(u + q, dim = 1).detach()
 
-    return (
+    out = (
         e -
         2 * (e @ rearrange(w, 'b d -> b d 1') @ rearrange(w, 'b d -> b 1 d')) +
         2 * (e @ rearrange(u, 'b d -> b d 1').detach() @ rearrange(q, 'b d -> b 1 d').detach())
     )
 
+    return rearrange(out, '... 1 -> ...')
+
 def rotate_to(src, tgt):
     # rotation trick STE (https://arxiv.org/abs/2410.06424) to get gradients through VQ layer.
     src, inverse = pack_one(src, '* d')
@@ -276,7 +293,7 @@ def rotate_to(src, tgt):
         safe_div(src, norm_src),
         safe_div(tgt, norm_tgt),
         src
-    ).squeeze()
+    )
 
     rotated = rotated_tgt * safe_div(norm_tgt, norm_src).detach()
 
@@ -511,7 +528,8 @@ class EuclideanCodebook(Module):
         mask = None,
         freeze_codebook = False,
         codebook_transform_fn: Callable | None = None,
-        ema_update_weight: Tensor | Callable | None = None
+        ema_update_weight: Tensor | Callable | None = None,
+        accum_ema_update = False
     ):
         needs_codebook_dim = x.ndim < 4
         sample_codebook_temp = default(sample_codebook_temp, self.sample_codebook_temp)
@@ -603,12 +621,16 @@ class EuclideanCodebook(Module):
             if callable(ema_update_weight):
                 ema_update_weight = ema_update_weight(embed_sum, cluster_size)
 
-            ema_inplace(self.cluster_size.data, cluster_size, self.decay, ema_update_weight)
-            ema_inplace(self.embed_avg.data, embed_sum, self.decay, ema_update_weight)
+            if accum_ema_update:
+                accum_grad_(self.cluster_size, cluster_size)
+                accum_grad_(self.embed_avg, embed_sum)
+            else:
+                ema_inplace(self.cluster_size, cluster_size, self.decay, ema_update_weight)
+                ema_inplace(self.embed_avg, embed_sum, self.decay, ema_update_weight)
 
-            if not self.manual_ema_update:
-                self.update_ema()
-                self.expire_codes_(x)
+                if not self.manual_ema_update:
+                    self.update_ema()
+                    self.expire_codes_(x)
 
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
@@ -743,7 +765,8 @@ class CosineSimCodebook(Module):
         mask = None,
         freeze_codebook = False,
         codebook_transform_fn: Callable | None = None,
-        ema_update_weight: Tensor | None = None
+        ema_update_weight: Tensor | None = None,
+        accum_ema_update = False
     ):
         needs_codebook_dim = x.ndim < 4
         sample_codebook_temp = default(sample_codebook_temp, self.sample_codebook_temp)
@@ -819,12 +842,17 @@ class CosineSimCodebook(Module):
             if callable(ema_update_weight):
                 ema_update_weight = ema_update_weight(embed_sum, bins)
 
-            ema_inplace(self.cluster_size.data, bins, self.decay, ema_update_weight)
-            ema_inplace(self.embed_avg.data, embed_sum, self.decay, ema_update_weight)
+            if accum_ema_update:
+                accum_grad_(self.cluster_size, bins)
+                accum_grad_(self.embed_avg, embed_sum)
+            else:
+
+                ema_inplace(self.cluster_size, bins, self.decay, ema_update_weight)
+                ema_inplace(self.embed_avg, embed_sum, self.decay, ema_update_weight)
 
-            if not self.manual_ema_update:
-                self.update_ema()
-                self.expire_codes_(x)
+                if not self.manual_ema_update:
+                    self.update_ema()
+                    self.expire_codes_(x)
 
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
@@ -870,7 +898,7 @@ class VectorQuantize(Module):
         stochastic_sample_codes = False,
         sample_codebook_temp = 1.,
         straight_through = False,
-        rotation_trick = True,  # Propagate grads through VQ layer w/ rotation trick: https://arxiv.org/abs/2410.06424 by @cfifty
+        rotation_trick = None,  # Propagate grads through VQ layer w/ rotation trick: https://arxiv.org/abs/2410.06424 by @cfifty
         sync_codebook = None,
         sync_affine_param = False,
         ema_update = True,
@@ -885,6 +913,8 @@ class VectorQuantize(Module):
         return_zeros_for_masked_padding = True
     ):
         super().__init__()
+        rotation_trick = default(rotation_trick, dim > 1) # only use rotation trick if feature dimension greater than 1
+
         self.dim = dim
         self.heads = heads
         self.separate_codebook_per_head = separate_codebook_per_head
@@ -1037,6 +1067,19 @@ class VectorQuantize(Module):
         if not exists(self.in_place_codebook_optimizer):
             return
 
+        # handle ddp, thanks to @gdoras
+
+        if self._codebook.use_ddp:
+
+            for param in self._codebook.parameters():
+                if not exists(param.grad):
+                    continue
+
+                distributed.all_reduce(param.grad)
+                param.grad /= distributed.get_world_size()
+
+        # optimizer step
+
         self.in_place_codebook_optimizer.step()
         self.in_place_codebook_optimizer.zero_grad()
 
@@ -1062,7 +1105,8 @@ class VectorQuantize(Module):
         freeze_codebook = None,
         return_loss_breakdown = False,
         codebook_transform_fn: Callable | None = None,
-        ema_update_weight: Tensor | None = None
+        ema_update_weight: Tensor | None = None,
+        accum_ema_update = False
     ):
         orig_input, input_requires_grad = x, x.requires_grad
 
@@ -1119,7 +1163,8 @@ class VectorQuantize(Module):
             mask = mask,
             freeze_codebook = freeze_codebook,
             codebook_transform_fn = codebook_transform_fn,
-            ema_update_weight = ema_update_weight
+            ema_update_weight = ema_update_weight,
+            accum_ema_update = accum_ema_update
         )
 
         # quantize
@@ -1222,7 +1267,7 @@ class VectorQuantize(Module):
             # calculate codebook diversity loss (negative of entropy) if needed
 
             if self.has_codebook_diversity_loss:
-                prob = (-distances * self.codebook_diversity_temperature).softmax(dim = -1)
+                prob = (distances * self.codebook_diversity_temperature).softmax(dim = -1)
                 avg_prob = reduce(prob, '... n l -> n l', 'mean')
                 codebook_diversity_loss = -entropy(avg_prob).mean()