vector-quantize-pytorch 1.18.7__tar.gz → 1.19.0__tar.gz

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.18.7
+Version: 1.19.0
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -144,6 +144,29 @@ quantized, indices, commit_loss = residual_vq(x)
 # (1, 1024, 256), (2, 1, 1024, 8), (2, 1, 8)
 ```
 
+
+<a href="https://arxiv.org/abs/2305.05065">This paper</a> out of Google Deepmind claims that residual vector quantization can induce hierarchical semantic ids for building a recommender system. In their scheme, they use increasing number of codes across depth for it to work. This repository supports that scheme as so
+
+```python
+import torch
+from vector_quantize_pytorch import ResidualVQ
+
+residual_vq = ResidualVQ(
+    dim = 2,
+    codebook_size = (5, 128, 256), # from top most hierarchy to lowest, 5 codes, 128 codes, then 256 codes
+)
+
+x = torch.randn(2, 2, 2)
+
+residual_vq.train()
+
+quantized, indices, commit_loss = residual_vq(x, freeze_codebook = True)
+
+quantized_out = residual_vq.get_output_from_indices(indices)
+
+assert torch.allclose(quantized, quantized_out, atol = 1e-5)
+```
+
 ## Initialization
 
 The SoundStream paper proposes that the codebook should be initialized by the kmeans centroids of the first batch. You can easily turn on this feature with one flag `kmeans_init = True`, for either `VectorQuantize` or `ResidualVQ` class
@@ -757,4 +780,15 @@ assert loss.item() >= 0
     volume  = {abs/2410.06424},
     url     = {https://api.semanticscholar.org/CorpusID:273229218}
 }
-```
+```
+
+```bibtex
+@article{Rajput2023RecommenderSW,
+    title   = {Recommender Systems with Generative Retrieval},
+    author  = {Shashank Rajput and Nikhil Mehta and Anima Singh and Raghunandan H. Keshavan and Trung Hieu Vu and Lukasz Heldt and Lichan Hong and Yi Tay and Vinh Q. Tran and Jonah Samost and Maciej Kula and Ed H. Chi and Maheswaran Sathiamoorthy},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2305.05065},
+    url     = {https://api.semanticscholar.org/CorpusID:258564854}
+}
+```

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/README.md

@@ -100,6 +100,29 @@ quantized, indices, commit_loss = residual_vq(x)
 # (1, 1024, 256), (2, 1, 1024, 8), (2, 1, 8)
 ```
 
+
+<a href="https://arxiv.org/abs/2305.05065">This paper</a> out of Google Deepmind claims that residual vector quantization can induce hierarchical semantic ids for building a recommender system. In their scheme, they use increasing number of codes across depth for it to work. This repository supports that scheme as so
+
+```python
+import torch
+from vector_quantize_pytorch import ResidualVQ
+
+residual_vq = ResidualVQ(
+    dim = 2,
+    codebook_size = (5, 128, 256), # from top most hierarchy to lowest, 5 codes, 128 codes, then 256 codes
+)
+
+x = torch.randn(2, 2, 2)
+
+residual_vq.train()
+
+quantized, indices, commit_loss = residual_vq(x, freeze_codebook = True)
+
+quantized_out = residual_vq.get_output_from_indices(indices)
+
+assert torch.allclose(quantized, quantized_out, atol = 1e-5)
+```
+
 ## Initialization
 
 The SoundStream paper proposes that the codebook should be initialized by the kmeans centroids of the first batch. You can easily turn on this feature with one flag `kmeans_init = True`, for either `VectorQuantize` or `ResidualVQ` class
@@ -713,4 +736,15 @@ assert loss.item() >= 0
     volume  = {abs/2410.06424},
     url     = {https://api.semanticscholar.org/CorpusID:273229218}
 }
-```
+```
+
+```bibtex
+@article{Rajput2023RecommenderSW,
+    title   = {Recommender Systems with Generative Retrieval},
+    author  = {Shashank Rajput and Nikhil Mehta and Anima Singh and Raghunandan H. Keshavan and Trung Hieu Vu and Lukasz Heldt and Lichan Hong and Yi Tay and Vinh Q. Tran and Jonah Samost and Maciej Kula and Ed H. Chi and Maheswaran Sathiamoorthy},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2305.05065},
+    url     = {https://api.semanticscholar.org/CorpusID:258564854}
+}
+```

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/pyproject.toml

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

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/tests/test_readme.py

@@ -201,6 +201,24 @@ def test_rq():
     x = torch.randn(1, 1024, 512)
     indices = quantizer(x)
 
+def test_tiger():
+    from vector_quantize_pytorch import ResidualVQ
+
+    residual_vq = ResidualVQ(
+        dim = 2,
+        codebook_size = (5, 128, 256),
+    )
+
+    x = torch.randn(2, 2, 2)
+
+    residual_vq.train()
+
+    quantized, indices, commit_loss = residual_vq(x, freeze_codebook = True)
+
+    quantized_out = residual_vq.get_output_from_indices(indices)  # pass your indices into here, but the indices must come during .eval(), as during training some of the indices are dropped out (-1)
+
+    assert torch.allclose(quantized, quantized_out, atol = 1e-5)
+
 def test_fsq():
     from vector_quantize_pytorch import FSQ
 

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/vector_quantize_pytorch/residual_fsq.py

@@ -185,7 +185,7 @@ class ResidualFSQ(Module):
             # check if seed is manually passed in
 
             if not exists(rand_quantize_dropout_fixed_seed):
-                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
+                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
 
             rand = random.Random(rand_quantize_dropout_fixed_seed)
 
@@ -296,7 +296,7 @@ class GroupedResidualFSQ(Module):
         x,
         return_all_codes = False
     ):
-        shape, split_dim = x.shape, self.split_dim
+        shape, split_dim, device = x.shape, self.split_dim, x.device
         assert shape[split_dim] == self.dim
 
         # split the feature dimension into groups
@@ -305,7 +305,7 @@ class GroupedResidualFSQ(Module):
 
         forward_kwargs = dict(
             return_all_codes = return_all_codes,
-            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
         )
 
         # invoke residual vq on each group

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/vector_quantize_pytorch/residual_lfq.py

@@ -31,8 +31,8 @@ def round_up_multiple(num, mult):
 def is_distributed():
     return dist.is_initialized() and dist.get_world_size() > 1
 
-def get_maybe_sync_seed(max_size = 10_000):
-    rand_int = torch.randint(0, max_size, ())
+def get_maybe_sync_seed(device, max_size = 10_000):
+    rand_int = torch.randint(0, max_size, (), device = device)
 
     if is_distributed():
         dist.all_reduce(rand_int)
@@ -162,7 +162,7 @@ class ResidualLFQ(Module):
             # check if seed is manually passed in
 
             if not exists(rand_quantize_dropout_fixed_seed):
-                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
+                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
 
             rand = random.Random(rand_quantize_dropout_fixed_seed)
 
@@ -262,7 +262,7 @@ class GroupedResidualLFQ(Module):
         mask = None,
         return_all_codes = False
     ):
-        shape, split_dim = x.shape, self.split_dim
+        shape, split_dim, device = x.shape, self.split_dim, x.device
         assert shape[split_dim] == self.dim
 
         # split the feature dimension into groups
@@ -272,7 +272,7 @@ class GroupedResidualLFQ(Module):
         forward_kwargs = dict(
             mask = mask,
             return_all_codes = return_all_codes,
-            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
         )
 
         # invoke residual vq on each group

{vector_quantize_pytorch-1.18.7 → vector_quantize_pytorch-1.19.0}/vector_quantize_pytorch/residual_vq.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-from typing import List
 
 import random
 from math import ceil
@@ -28,6 +27,12 @@ def first(it):
 def default(val, d):
     return val if exists(val) else d
 
+def cast_tuple(t, length = 1):
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+def unique(arr):
+    return list({*arr})
+
 def round_up_multiple(num, mult):
     return ceil(num / mult) * mult
 
@@ -36,8 +41,8 @@ def round_up_multiple(num, mult):
 def is_distributed():
     return dist.is_initialized() and dist.get_world_size() > 1
 
-def get_maybe_sync_seed(max_size = 10_000):
-    rand_int = torch.randint(0, max_size, ())
+def get_maybe_sync_seed(device, max_size = 10_000):
+    rand_int = torch.randint(0, max_size, (), device = device)
 
     if is_distributed():
         dist.all_reduce(rand_int)
@@ -110,7 +115,8 @@ class ResidualVQ(Module):
         self,
         *,
         dim,
-        num_quantizers,
+        num_quantizers: int | None = None,
+        codebook_size: int | tuple[int, ...],
         codebook_dim = None,
         shared_codebook = False,
         heads = 1,
@@ -124,6 +130,8 @@ class ResidualVQ(Module):
     ):
         super().__init__()
         assert heads == 1, 'residual vq is not compatible with multi-headed codes'
+        assert exists(num_quantizers) or isinstance(codebook_size, tuple)
+
         codebook_dim = default(codebook_dim, dim)
         codebook_input_dim = codebook_dim * heads
 
@@ -132,8 +140,6 @@ class ResidualVQ(Module):
         self.project_out = nn.Linear(codebook_input_dim, dim) if requires_projection else nn.Identity()
         self.has_projections = requires_projection
 
-        self.num_quantizers = num_quantizers
-
         self.accept_image_fmap = accept_image_fmap
 
         self.implicit_neural_codebook = implicit_neural_codebook
@@ -150,7 +156,21 @@ class ResidualVQ(Module):
                 manual_in_place_optimizer_update = True
             )
 
-        self.layers = ModuleList([VectorQuantize(dim = codebook_dim, codebook_dim = codebook_dim, accept_image_fmap = accept_image_fmap, **vq_kwargs) for _ in range(num_quantizers)])
+        # take care of maybe different codebook sizes across depth, used in TIGER paper https://arxiv.org/abs/2305.05065
+
+        codebook_sizes = cast_tuple(codebook_size, num_quantizers)
+
+        num_quantizers = len(codebook_sizes)
+        self.num_quantizers = num_quantizers
+
+        assert len(codebook_sizes) == num_quantizers
+
+        self.codebook_sizes = codebook_sizes
+        self.uniform_codebook_size = len(unique(codebook_sizes)) == 1
+
+        # define vq across layers
+
+        self.layers = ModuleList([VectorQuantize(dim = codebook_dim, codebook_size = layer_codebook_size, codebook_dim = codebook_dim, accept_image_fmap = accept_image_fmap, **vq_kwargs) for layer_codebook_size in codebook_sizes])
 
         assert all([not vq.has_projections for vq in self.layers])
 
@@ -167,6 +187,8 @@ class ResidualVQ(Module):
 
         if implicit_neural_codebook:
             self.mlps = ModuleList([MLP(dim = codebook_dim, l2norm_output = first(self.layers).use_cosine_sim, **mlp_kwargs) for _ in range(num_quantizers - 1)])
+        else:
+            self.mlps = (None,) * (num_quantizers - 1)
 
         # sharing codebook logic
 
@@ -175,6 +197,8 @@ class ResidualVQ(Module):
         if not shared_codebook:
             return
 
+        assert self.uniform_codebook_size
+
         first_vq, *rest_vq = self.layers
         codebook = first_vq._codebook
 
@@ -192,8 +216,13 @@ class ResidualVQ(Module):
     @property
     def codebooks(self):
         codebooks = [layer._codebook.embed for layer in self.layers]
-        codebooks = torch.stack(codebooks, dim = 0)
-        codebooks = rearrange(codebooks, 'q 1 c d -> q c d')
+
+        codebooks = tuple(rearrange(codebook, '1 ... -> ...') for codebook in codebooks)
+
+        if not self.uniform_codebook_size:
+            return codebooks
+
+        codebooks = torch.stack(codebooks)
         return codebooks
 
     def get_codes_from_indices(self, indices):
@@ -216,13 +245,12 @@ class ResidualVQ(Module):
         mask = indices == -1.
         indices = indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
 
-        if not self.implicit_neural_codebook:
-            # gather all the codes
+        if not self.implicit_neural_codebook and self.uniform_codebook_size:
 
             all_codes = get_at('q [c] d, b n q -> q b n d', self.codebooks, indices)
 
         else:
-            # else if using implicit neural codebook, codes will need to be derived layer by layer
+            # else if using implicit neural codebook, or non uniform codebook sizes, codes will need to be derived layer by layer
 
             code_transform_mlps = (None, *self.mlps)
 
@@ -261,7 +289,7 @@ class ResidualVQ(Module):
         self,
         x,
         mask = None,
-        indices: Tensor | List[Tensor] | None = None,
+        indices: Tensor | list[Tensor] | None = None,
         return_all_codes = False,
         sample_codebook_temp = None,
         freeze_codebook = False,
@@ -296,7 +324,7 @@ class ResidualVQ(Module):
             # check if seed is manually passed in
 
             if not exists(rand_quantize_dropout_fixed_seed):
-                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
+                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
 
             rand = random.Random(rand_quantize_dropout_fixed_seed)
 
@@ -452,7 +480,7 @@ class GroupedResidualVQ(Module):
         freeze_codebook = False,
         mask = None,
     ):
-        shape, split_dim = x.shape, self.split_dim
+        shape, split_dim, device = x.shape, self.split_dim, x.device
         assert shape[split_dim] == self.dim
 
         # split the feature dimension into groups
@@ -468,7 +496,7 @@ class GroupedResidualVQ(Module):
             sample_codebook_temp = sample_codebook_temp,
             mask = mask,
             freeze_codebook = freeze_codebook,
-            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
         )
 
         # invoke residual vq on each group