vector-quantize-pytorch 1.18.3__tar.gz → 1.18.6__tar.gz

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.18.3
+Version: 1.18.6
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -655,16 +655,6 @@ assert loss.item() >= 0
 }
 ```
 
-```bibtex
-@article{Liu2023BridgingDA,
-    title   = {Bridging Discrete and Backpropagation: Straight-Through and Beyond},
-    author  = {Liyuan Liu and Chengyu Dong and Xiaodong Liu and Bin Yu and Jianfeng Gao},
-    journal = {ArXiv},
-    year    = {2023},
-    volume  = {abs/2304.08612}
-}
-```
-
 ```bibtex
 @inproceedings{huh2023improvedvqste,
     title   = {Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks},

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/README.md

@@ -611,16 +611,6 @@ assert loss.item() >= 0
 }
 ```
 
-```bibtex
-@article{Liu2023BridgingDA,
-    title   = {Bridging Discrete and Backpropagation: Straight-Through and Beyond},
-    author  = {Liyuan Liu and Chengyu Dong and Xiaodong Liu and Bin Yu and Jianfeng Gao},
-    journal = {ArXiv},
-    year    = {2023},
-    volume  = {abs/2304.08612}
-}
-```
-
 ```bibtex
 @inproceedings{huh2023improvedvqste,
     title   = {Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks},

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/pyproject.toml

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

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/vector_quantize_pytorch/lookup_free_quantization.py

@@ -103,7 +103,6 @@ class LFQ(Module):
         commitment_loss_weight = 0.,
         diversity_gamma = 1.,
         straight_through_activation = nn.Identity(),
-        scale_trick = False,                        # @cfifty Fifty et al. https://arxiv.org/abs/2410.06424
         num_codebooks = 1,
         keep_num_codebooks_dim = None,
         codebook_scale = 1.,                        # for residual LFQ, codebook scaled down by 2x at each layer
@@ -161,9 +160,6 @@ class LFQ(Module):
 
         self.activation = straight_through_activation
 
-        assert not (scale_trick and spherical)
-        self.scale_trick = scale_trick
-
         # whether to use BSQ (binary spherical quantization)
 
         self.spherical = spherical
@@ -326,12 +322,7 @@ class LFQ(Module):
 
             if self.training:
                 x = self.activation(x)
-
-                if self.scale_trick:
-                    x = x * (quantized / x).detach()
-                else:
-                    x = x + (quantized - x).detach()
-
+                x = x + (quantized - x).detach()
             else:
                 x = quantized
 

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/vector_quantize_pytorch/residual_fsq.py

@@ -1,6 +1,6 @@
 import random
 from math import log2
-from functools import partial
+from functools import partial, cache
 
 from typing import List
 
@@ -9,6 +9,7 @@ from torch import nn
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
 from torch.amp import autocast
+import torch.distributed as dist
 
 from vector_quantize_pytorch.finite_scalar_quantization import FSQ
 
@@ -30,6 +31,12 @@ def default(val, d):
 def round_up_multiple(num, mult):
     return ceil(num / mult) * mult
 
+# distributed helpers
+
+@cache
+def is_distributed():
+    return dist.is_initialized() and dist.get_world_size() > 1
+
 # main class
 
 class ResidualFSQ(Module):
@@ -38,9 +45,10 @@ class ResidualFSQ(Module):
     def __init__(
         self,
         *,
-        dim,
         levels: List[int],
         num_quantizers,
+        dim = None,
+        is_channel_first = False,
         quantize_dropout = False,
         quantize_dropout_cutoff_index = 0,
         quantize_dropout_multiple_of = 1,
@@ -48,12 +56,14 @@ class ResidualFSQ(Module):
     ):
         super().__init__()
         codebook_dim = len(levels)
+        dim = default(dim, codebook_dim)
 
         requires_projection = codebook_dim != dim
         self.project_in = nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity()
         self.project_out = nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
         self.has_projections = requires_projection
 
+        self.is_channel_first = is_channel_first
         self.num_quantizers = num_quantizers
 
         self.levels = levels
@@ -143,10 +153,18 @@ class ResidualFSQ(Module):
     ):
         num_quant, quant_dropout_multiple_of, device = self.num_quantizers, self.quantize_dropout_multiple_of, x.device
 
+        # handle channel first
+
+        if self.is_channel_first:
+            x = rearrange(x, 'b d ... -> b ... d')
+            x, ps = pack([x], 'b * d')
+
+        # maybe project in
+
         x = self.project_in(x)
 
         quantized_out = 0.
-        residual = first(self.layers).bound(x)
+        residual = x
 
         all_indices = []
 
@@ -156,7 +174,19 @@ class ResidualFSQ(Module):
         # also prepare null indices
 
         if should_quantize_dropout:
-            rand = random.Random(rand_quantize_dropout_fixed_seed) if exists(rand_quantize_dropout_fixed_seed) else random
+
+            if exists(rand_quantize_dropout_fixed_seed):
+                # seed is manually passed in
+                rand = random.Random(rand_quantize_dropout_fixed_seed)
+
+            elif is_distributed():
+                # in distributed environment, synchronize a random seed value if not given
+                t = torch.tensor(random.randrange(10_000), device = device)
+                dropout_seed = dist.all_reduce(t).item()
+                rand = random.Random(dropout_seed)
+
+            else:
+                rand = random
 
             rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
 
@@ -175,6 +205,7 @@ class ResidualFSQ(Module):
                     continue
 
                 quantized, indices = layer(residual / scale)
+
                 quantized = quantized * scale
 
                 residual = residual - quantized.detach()
@@ -190,6 +221,17 @@ class ResidualFSQ(Module):
 
         all_indices = torch.stack(all_indices, dim = -1)
 
+        # channel first out
+
+        if self.is_channel_first:
+            quantized_out, = unpack(quantized_out, ps, 'b * d')
+            all_indices, = unpack(all_indices, ps, 'b * d')
+
+            quantized_out = rearrange(quantized_out, 'b ... d -> b d ...')
+            all_indices = rearrange(all_indices, 'b ... d -> b d ...')
+
+        # return
+
         ret = (quantized_out, all_indices)
 
         if not return_all_codes:

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/vector_quantize_pytorch/residual_lfq.py

@@ -1,12 +1,13 @@
 import random
 from math import log2
-from functools import partial
+from functools import partial, cache
 
 import torch
 from torch import nn
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
 from torch.amp import autocast
+import torch.distributed as dist
 
 from vector_quantize_pytorch.lookup_free_quantization import LFQ
 
@@ -25,6 +26,12 @@ def default(val, d):
 def round_up_multiple(num, mult):
     return ceil(num / mult) * mult
 
+# distributed helpers
+
+@cache
+def is_distributed():
+    return dist.is_initialized() and dist.get_world_size() > 1
+
 # main class
 
 class ResidualLFQ(Module):
@@ -144,7 +151,19 @@ class ResidualLFQ(Module):
         # also prepare null indices and loss
 
         if should_quantize_dropout:
-            rand = random.Random(rand_quantize_dropout_fixed_seed) if exists(rand_quantize_dropout_fixed_seed) else random
+
+            if exists(rand_quantize_dropout_fixed_seed):
+                # seed is manually passed in
+                rand = random.Random(rand_quantize_dropout_fixed_seed)
+
+            elif is_distributed():
+                # in distributed environment, synchronize a random seed value if not given
+                t = torch.tensor(random.randrange(10_000), device = device)
+                dropout_seed = dist.all_reduce(t).item()
+                rand = random.Random(dropout_seed)
+
+            else:
+                rand = random
 
             rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
 

{vector_quantize_pytorch-1.18.3 → vector_quantize_pytorch-1.18.6}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -103,7 +103,6 @@ def gumbel_sample(
     temperature = 1.,
     stochastic = False,
     straight_through = False,
-    reinmax = False,
     dim = -1,
     training = True
 ):
@@ -117,23 +116,11 @@ def gumbel_sample(
     ind = sampling_logits.argmax(dim = dim)
     one_hot = F.one_hot(ind, size).type(dtype)
 
-    assert not (reinmax and not straight_through), 'reinmax can only be turned on if using straight through gumbel softmax'
-
     if not straight_through or temperature <= 0. or not training:
         return ind, one_hot
 
-    # use reinmax for better second-order accuracy - https://arxiv.org/abs/2304.08612
-    # algorithm 2
-
-    if reinmax:
-        π0 = logits.softmax(dim = dim)
-        π1 = (one_hot + (logits / temperature).softmax(dim = dim)) / 2
-        π1 = ((log(π1) - logits).detach() + logits).softmax(dim = 1)
-        π2 = 2 * π1 - 0.5 * π0
-        one_hot = π2 - π2.detach() + one_hot
-    else:
-        π1 = (logits / temperature).softmax(dim = dim)
-        one_hot = one_hot + π1 - π1.detach()
+    π1 = (logits / temperature).softmax(dim = dim)
+    one_hot = one_hot + π1 - π1.detach()
 
     return ind, one_hot
 
@@ -828,7 +815,6 @@ class VectorQuantize(Module):
         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
-        reinmax = False,  # using reinmax for improved straight-through, assuming straight through helps at all
         sync_codebook = None,
         sync_affine_param = False,
         ema_update = True,
@@ -895,7 +881,6 @@ class VectorQuantize(Module):
         gumbel_sample_fn = partial(
             gumbel_sample,
             stochastic = stochastic_sample_codes,
-            reinmax = reinmax,
             straight_through = straight_through
         )