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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.20.0
+Version: 1.20.2
 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.0 → vector_quantize_pytorch-1.20.2}/pyproject.toml

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

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

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import Callable
 
 import torch
@@ -38,10 +39,11 @@ class SimVQ(Module):
         self,
         dim,
         codebook_size,
+        codebook_transform: Module | None = None,
         init_fn: Callable = identity,
         accept_image_fmap = False,
-        rotation_trick = True,  # works even better with rotation trick turned on, with no asymmetric commit loss or straight through
-        commit_loss_input_to_quantize_weight = 0.25,
+        rotation_trick = True,  # works even better with rotation trick turned on, with no straight through and the commit loss from input to quantize
+        input_to_quantize_commit_loss_weight = 0.25,
     ):
         super().__init__()
         self.accept_image_fmap = accept_image_fmap
@@ -51,7 +53,11 @@ class SimVQ(Module):
 
         # the codebook is actually implicit from a linear layer from frozen gaussian or uniform
 
-        self.codebook_to_codes = nn.Linear(dim, dim, bias = False)
+        if not exists(codebook_transform):
+            codebook_transform = nn.Linear(dim, dim, bias = False)
+
+        self.codebook_to_codes = codebook_transform
+
         self.register_buffer('codebook', codebook)
 
 
@@ -59,11 +65,10 @@ class SimVQ(Module):
         # https://arxiv.org/abs/2410.06424
 
         self.rotation_trick = rotation_trick
-        self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
         # commit loss weighting - weighing input to quantize a bit less is crucial for it to work
 
-        self.commit_loss_input_to_quantize_weight = commit_loss_input_to_quantize_weight
+        self.input_to_quantize_commit_loss_weight = input_to_quantize_commit_loss_weight
 
     def forward(
         self,
@@ -83,18 +88,18 @@ class SimVQ(Module):
 
         quantized = get_at('[c] d, b n -> b n d', implicit_codebook, indices)
 
+        # commit loss and straight through, as was done in the paper
+
+        commit_loss = F.mse_loss(x.detach(), quantized)
+
         if self.rotation_trick:
             # rotation trick from @cfifty
-
             quantized = rotate_from_to(quantized, x)
-
-            commit_loss = self.zero
         else:
-            # commit loss and straight through, as was done in the paper
 
             commit_loss = (
-                F.mse_loss(x, quantized.detach()) * self.commit_loss_input_to_quantize_weight +
-                F.mse_loss(x.detach(), quantized)
+                commit_loss + 
+                F.mse_loss(x, quantized.detach()) * self.input_to_quantize_commit_loss_weight
             )
 
             quantized = (quantized - x).detach() + x
@@ -115,6 +120,11 @@ if __name__ == '__main__':
 
     sim_vq = SimVQ(
         dim = 512,
+        codebook_transform = nn.Sequential(
+            nn.Linear(512, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, 512)
+        ),
         codebook_size = 1024,
         accept_image_fmap = True
     )