Vector quantization implementation. Currently supports only euclidean distance. Args: dim (int): Dimension codebook_size (int): Codebook size codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim. decay (float): Decay for
| 232 | |
| 233 | |
| 234 | class VectorQuantization(nn.Module): |
| 235 | """Vector quantization implementation. |
| 236 | Currently supports only euclidean distance. |
| 237 | Args: |
| 238 | dim (int): Dimension |
| 239 | codebook_size (int): Codebook size |
| 240 | codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim. |
| 241 | decay (float): Decay for exponential moving average over the codebooks. |
| 242 | epsilon (float): Epsilon value for numerical stability. |
| 243 | kmeans_init (bool): Whether to use kmeans to initialize the codebooks. |
| 244 | kmeans_iters (int): Number of iterations used for kmeans initialization. |
| 245 | threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes |
| 246 | that have an exponential moving average cluster size less than the specified threshold with |
| 247 | randomly selected vector from the current batch. |
| 248 | commitment_weight (float): Weight for commitment loss. |
| 249 | """ |
| 250 | |
| 251 | def __init__( |
| 252 | self, |
| 253 | dim: int, |
| 254 | codebook_size: int, |
| 255 | codebook_dim: tp.Optional[int] = None, |
| 256 | decay: float = 0.99, |
| 257 | epsilon: float = 1e-5, |
| 258 | kmeans_init: bool = True, |
| 259 | kmeans_iters: int = 50, |
| 260 | threshold_ema_dead_code: int = 2, |
| 261 | commitment_weight: float = 1.0, |
| 262 | ): |
| 263 | super().__init__() |
| 264 | _codebook_dim: int = default(codebook_dim, dim) |
| 265 | |
| 266 | requires_projection = _codebook_dim != dim |
| 267 | self.project_in = ( |
| 268 | nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity() |
| 269 | ) |
| 270 | self.project_out = ( |
| 271 | nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity() |
| 272 | ) |
| 273 | |
| 274 | self.epsilon = epsilon |
| 275 | self.commitment_weight = commitment_weight |
| 276 | |
| 277 | self._codebook = EuclideanCodebook( |
| 278 | dim=_codebook_dim, |
| 279 | codebook_size=codebook_size, |
| 280 | kmeans_init=kmeans_init, |
| 281 | kmeans_iters=kmeans_iters, |
| 282 | decay=decay, |
| 283 | epsilon=epsilon, |
| 284 | threshold_ema_dead_code=threshold_ema_dead_code, |
| 285 | ) |
| 286 | self.codebook_size = codebook_size |
| 287 | |
| 288 | @property |
| 289 | def codebook(self): |
| 290 | return self._codebook.embed |
| 291 |