vector-quantize-pytorch 1.20.2__tar.gz → 1.20.4__tar.gz

{vector_quantize_pytorch-1.20.2 → vector_quantize_pytorch-1.20.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.20.2
+Version: 1.20.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.20.2 → vector_quantize_pytorch-1.20.4}/examples/autoencoder_sim_vq.py

@@ -14,7 +14,10 @@ lr = 3e-4
 train_iter = 10000
 num_codes = 256
 seed = 1234
-rotation_trick = True
+
+rotation_trick = True # rotation trick instead ot straight-through
+use_mlp = True        # use a one layer mlp with relu instead of linear
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def SimVQAutoEncoder(**vq_kwargs):
@@ -77,7 +80,12 @@ torch.random.manual_seed(seed)
 
 model = SimVQAutoEncoder(
     codebook_size = num_codes,
-    rotation_trick = rotation_trick
+    rotation_trick = rotation_trick,
+    codebook_transform = nn.Sequential(
+        nn.Linear(32, 128),
+        nn.ReLU(),
+        nn.Linear(128, 32),
+    ) if use_mlp else None
 ).to(device)
 
 opt = torch.optim.AdamW(model.parameters(), lr=lr)

{vector_quantize_pytorch-1.20.2 → vector_quantize_pytorch-1.20.4}/pyproject.toml

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

{vector_quantize_pytorch-1.20.2 → vector_quantize_pytorch-1.20.4}/tests/test_readme.py

@@ -362,3 +362,17 @@ def test_latent_q():
 
     assert image_feats.shape == quantized.shape
     assert (quantized == quantizer.indices_to_codes(indices)).all()
+
+def test_sim_vq():
+    from vector_quantize_pytorch import SimVQ
+
+    sim_vq = SimVQ(
+        dim = 512,
+        codebook_size = 1024,
+    )
+
+    x = torch.randn(1, 1024, 512)
+    quantized, indices, commit_loss = sim_vq(x)
+
+    assert x.shape == quantized.shape
+    assert torch.allclose(quantized, sim_vq.indices_to_codes(indices), atol = 1e-6)

{vector_quantize_pytorch-1.20.2 → vector_quantize_pytorch-1.20.4}/vector_quantize_pytorch/sim_vq.py

@@ -9,7 +9,7 @@ import torch.nn.functional as F
 from einx import get_at
 from einops import rearrange, pack, unpack
 
-from vector_quantize_pytorch.vector_quantize_pytorch import rotate_from_to
+from vector_quantize_pytorch.vector_quantize_pytorch import rotate_to
 
 # helper functions
 
@@ -56,9 +56,9 @@ class SimVQ(Module):
         if not exists(codebook_transform):
             codebook_transform = nn.Linear(dim, dim, bias = False)
 
-        self.codebook_to_codes = codebook_transform
+        self.code_transform = codebook_transform
 
-        self.register_buffer('codebook', codebook)
+        self.register_buffer('frozen_codebook', codebook)
 
 
         # whether to use rotation trick from Fifty et al. 
@@ -70,6 +70,24 @@ class SimVQ(Module):
 
         self.input_to_quantize_commit_loss_weight = input_to_quantize_commit_loss_weight
 
+    @property
+    def codebook(self):
+        return self.code_transform(self.frozen_codebook)
+
+    def indices_to_codes(
+        self,
+        indices
+    ):
+        implicit_codebook = self.codebook
+
+        frozen_codes = get_at('[c] d, b ... -> b ... d', self.frozen_codebook, indices)
+        quantized = self.code_transform(frozen_codes)
+
+        if self.accept_image_fmap:
+            quantized = rearrange(quantized, 'b ... d -> b d ...')
+
+        return quantized
+
     def forward(
         self,
         x
@@ -78,7 +96,7 @@ class SimVQ(Module):
             x = rearrange(x, 'b d h w -> b h w d')
             x, inverse_pack = pack_one(x, 'b * d')
 
-        implicit_codebook = self.codebook_to_codes(self.codebook)
+        implicit_codebook = self.codebook
 
         with torch.no_grad():
             dist = torch.cdist(x, implicit_codebook)
@@ -94,7 +112,7 @@ class SimVQ(Module):
 
         if self.rotation_trick:
             # rotation trick from @cfifty
-            quantized = rotate_from_to(quantized, x)
+            quantized = rotate_to(x, quantized)
         else:
 
             commit_loss = (

{vector_quantize_pytorch-1.20.2 → vector_quantize_pytorch-1.20.4}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -250,21 +250,21 @@ def efficient_rotation_trick_transform(u, q, e):
         2 * (e @ rearrange(u, 'b d -> b d 1').detach() @ rearrange(q, 'b d -> b 1 d').detach())
     )
 
-def rotate_from_to(src, tgt):
+def rotate_to(src, tgt):
     # rotation trick STE (https://arxiv.org/abs/2410.06424) to get gradients through VQ layer.
-    tgt, inverse = pack_one(tgt, '* d')
-    src, _ = pack_one(src, '* d')
+    src, inverse = pack_one(src, '* d')
+    tgt, _ = pack_one(tgt, '* d')
 
-    norm_tgt = tgt.norm(dim = -1, keepdim = True)
     norm_src = src.norm(dim = -1, keepdim = True)
+    norm_tgt = tgt.norm(dim = -1, keepdim = True)
 
-    rotated_src = efficient_rotation_trick_transform(
-        safe_div(tgt, norm_tgt),
+    rotated_tgt = efficient_rotation_trick_transform(
         safe_div(src, norm_src),
-        tgt
+        safe_div(tgt, norm_tgt),
+        src
     ).squeeze()
 
-    rotated = rotated_src * safe_div(norm_src, norm_tgt).detach()
+    rotated = rotated_tgt * safe_div(norm_tgt, norm_src).detach()
 
     return inverse(rotated)
 
@@ -1118,7 +1118,7 @@ class VectorQuantize(Module):
             commit_quantize = maybe_detach(quantize)
 
             if self.rotation_trick:
-                quantize = rotate_from_to(quantize, x)
+                quantize = rotate_to(x, quantize)
             else:
                 # standard STE to get gradients through VQ layer.
                 quantize = x + (quantize - x).detach()