Add code/cube3d/model/autoencoder/spherical_vq.py

code/cube3d/model/autoencoder/spherical_vq.py

@@ -0,0 +1,159 @@
+import sys
+from typing import Literal, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cube3d.model.transformers.norm import RMSNorm
+
+
+class SphericalVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_codes: int,
+        width: Optional[int] = None,
+        codebook_regularization: Literal["batch_norm", "kl"] = "batch_norm",
+    ):
+        """
+        Initializes the SphericalVQ module.
+        Args:
+            embed_dim (int): The dimensionality of the embeddings.
+            num_codes (int): The number of codes in the codebook.
+            width (Optional[int], optional): The width of the input. Defaults to None.
+        Raises:
+            ValueError: If beta is not in the range [0, 1].
+        """
+        super().__init__()
+
+        self.num_codes = num_codes
+
+        self.codebook = nn.Embedding(num_codes, embed_dim)
+        self.codebook.weight.data.uniform_(-1.0 / num_codes, 1.0 / num_codes)
+
+        width = width or embed_dim
+        if width != embed_dim:
+            self.c_in = nn.Linear(width, embed_dim)
+            self.c_x = nn.Linear(width, embed_dim)  # shortcut
+            self.c_out = nn.Linear(embed_dim, width)
+        else:
+            self.c_in = self.c_out = self.c_x = nn.Identity()
+
+        self.norm = RMSNorm(embed_dim, elementwise_affine=False)
+        self.cb_reg = codebook_regularization
+        if self.cb_reg == "batch_norm":
+            self.cb_norm = nn.BatchNorm1d(embed_dim, track_running_stats=False)
+        else:
+            self.cb_weight = nn.Parameter(torch.ones([embed_dim]))
+            self.cb_bias = nn.Parameter(torch.zeros([embed_dim]))
+            self.cb_norm = lambda x: x.mul(self.cb_weight).add_(self.cb_bias)
+
+    def get_codebook(self):
+        """
+        Retrieves the normalized codebook weights.
+        This method applies a series of normalization operations to the
+        codebook weights, ensuring they are properly scaled and normalized
+        before being returned.
+        Returns:
+            torch.Tensor: The normalized weights of the codebook.
+        """
+
+        return self.norm(self.cb_norm(self.codebook.weight))
+
+    @torch.no_grad()
+
+    def lookup_codebook(self, q: torch.Tensor):
+        """
+        Perform a lookup in the codebook and process the result.
+        This method takes an input tensor of indices, retrieves the corresponding
+        embeddings from the codebook, and applies a transformation to the retrieved
+        embeddings.
+        Args:
+            q (torch.Tensor): A tensor containing indices to look up in the codebook.
+        Returns:
+            torch.Tensor: The transformed embeddings retrieved from the codebook.
+        """
+
+        # normalize codebook
+        z_q = F.embedding(q, self.get_codebook())
+        z_q = self.c_out(z_q)
+        return z_q
+
+    @torch.no_grad()
+    def lookup_codebook_latents(self, q: torch.Tensor):
+        """
+        Retrieves the latent representations from the codebook corresponding to the given indices.
+        Args:
+            q (torch.Tensor): A tensor containing the indices of the codebook entries to retrieve.
+                              The indices should be integers and correspond to the rows in the codebook.
+        Returns:
+            torch.Tensor: A tensor containing the latent representations retrieved from the codebook.
+                          The shape of the returned tensor depends on the shape of the input indices
+                          and the dimensionality of the codebook entries.
+        """
+
+        # normalize codebook
+        z_q = F.embedding(q, self.get_codebook())
+        return z_q
+
+    def quantize(self, z: torch.Tensor):
+        """
+        Quantizes the latent codes z with the codebook
+
+        Args:
+                z (Tensor): B x ... x F
+        """
+
+        # normalize codebook
+        codebook = self.get_codebook()
+        # the process of finding quantized codes is non differentiable
+        with torch.no_grad():
+            # flatten z
+            z_flat = z.view(-1, z.shape[-1])
+
+            # calculate distance and find the closest code
+            d = torch.cdist(z_flat, codebook)
+            q = torch.argmin(d, dim=1)  # num_ele
+
+        z_q = codebook[q, :].reshape(*z.shape[:-1], -1)
+        q = q.view(*z.shape[:-1])
+
+        return z_q, {"z": z.detach(), "q": q}
+
+    def straight_through_approximation(self, z, z_q):
+        """passed gradient from z_q to z"""
+        z_q = z + (z_q - z).detach()
+        return z_q
+
+    def forward(self, z: torch.Tensor):
+        """
+        Forward pass of the spherical vector quantization autoencoder.
+        Args:
+            z (torch.Tensor): Input tensor of shape (batch_size, ..., feature_dim).
+        Returns:
+            Tuple[torch.Tensor, Dict[str, Any]]:
+                - z_q (torch.Tensor): The quantized output tensor after applying the
+                  straight-through approximation and output projection.
+                - ret_dict (Dict[str, Any]): A dictionary containing additional
+                  information:
+                    - "z_q" (torch.Tensor): Detached quantized tensor.
+                    - "q" (torch.Tensor): Indices of the quantized vectors.
+                    - "perplexity" (torch.Tensor): The perplexity of the quantization,
+                      calculated as the exponential of the negative sum of the
+                      probabilities' log values.
+        """
+
+        with torch.autocast(device_type=z.device.type, enabled=False):
+            # work in full precision
+            z = z.float()
+
+            # project and normalize
+            z_e = self.norm(self.c_in(z))
+            z_q, ret_dict = self.quantize(z_e)
+
+            ret_dict["z_q"] = z_q.detach()
+            z_q = self.straight_through_approximation(z_e, z_q)
+            z_q = self.c_out(z_q)
+
+        return z_q, ret_dict