vector-quantize-pytorch 1.31.0__tar.gz → 1.31.2__tar.gz

{vector_quantize_pytorch-1.31.0 → vector_quantize_pytorch-1.31.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: vector-quantize-pytorch
-Version: 1.31.0
+Version: 1.31.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
@@ -36,6 +36,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
+Requires-Dist: torch-einops-utils>=0.1.12
 Requires-Dist: torch>=2.4
 Provides-Extra: examples
 Requires-Dist: torchvision; extra == 'examples'

{vector_quantize_pytorch-1.31.0 → vector_quantize_pytorch-1.31.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.31.0"
+version = "1.31.2"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -23,9 +23,10 @@ classifiers=[
   ]
 
 dependencies = [
-    "torch>=2.4",
     "einops>=0.8.0",
     "einx>=0.3.0",
+    "torch>=2.4",
+    "torch-einops-utils>=0.1.12",
 ]
 
 [project.urls]

{vector_quantize_pytorch-1.31.0 → vector_quantize_pytorch-1.31.2}/vector_quantize_pytorch/__init__.py

@@ -4,6 +4,7 @@ from vector_quantize_pytorch.random_projection_quantizer import RandomProjection
 from vector_quantize_pytorch.finite_scalar_quantization import FSQ
 from vector_quantize_pytorch.finite_scalar_perturbation import FSP
 from vector_quantize_pytorch.lookup_free_quantization import LFQ
+from vector_quantize_pytorch.evo_lookup_free_quantization import EvoLFQ
 from vector_quantize_pytorch.residual_lfq import ResidualLFQ, GroupedResidualLFQ
 from vector_quantize_pytorch.residual_fsq import ResidualFSQ, GroupedResidualFSQ
 from vector_quantize_pytorch.latent_quantization import LatentQuantize

vector_quantize_pytorch-1.31.2/vector_quantize_pytorch/evo_lookup_free_quantization.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+from collections import namedtuple
+
+import torch
+from torch import nn, Tensor, is_tensor
+from torch.nn import Module
+import torch.nn.functional as F
+
+from einops import rearrange, repeat, reduce
+from torch_einops_utils import pack_with_inverse, temp_eval
+
+from vector_quantize_pytorch.lookup_free_quantization import LFQ
+
+# constants
+
+Return = namedtuple('Return', ['reconstructed', 'indices', 'entropy_aux_loss'])
+Result = namedtuple('Result', ['pop_bits', 'best_gene', 'best_fitness', 'best_decoded'])
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+# main class
+
+class EvoLFQ(Module):
+    def __init__(
+        self,
+        encoder: Module,
+        decoder: Module,
+        lfq: LFQ | None = None,
+        *,
+        dim: int | None = None,
+        codebook_size: int | None = None,
+        num_codebooks: int = 1,
+        pop_size: int = 64,
+        mutation_rate: float = 0.02,
+        tournament_size: int = 2,
+        elitism_count: int = 1,
+        generations: int = 50,
+        **lfq_kwargs
+    ):
+        super().__init__()
+        assert pop_size > elitism_count, 'pop_size must be greater than elitism_count'
+
+        self.encoder = encoder
+        self.decoder = decoder
+
+        if not exists(lfq):
+            assert exists(dim) or exists(codebook_size), 'either lfq instance or dim / codebook_size must be supplied to EvoLFQ'
+            lfq = LFQ(
+                dim = dim,
+                codebook_size = codebook_size,
+                num_codebooks = num_codebooks,
+                **lfq_kwargs
+            )
+
+        self.lfq = lfq
+        self.pop_size = pop_size
+        self.mutation_rate = mutation_rate
+        self.tournament_size = tournament_size
+        self.elitism_count = elitism_count
+        self.generations = generations
+
+        self.register_buffer('zero', torch.tensor(0.), persistent = False)
+
+    @property
+    def device(self):
+        return self.zero.device
+
+    def forward(self, x, **kwargs):
+        latents = self.encoder(x)
+        is_2d = latents.ndim == 2
+
+        if is_2d:
+            latents = rearrange(latents, 'b d -> b 1 d')
+
+        quantized, indices, aux_loss = self.lfq(latents, **kwargs)
+
+        if is_2d:
+            quantized = rearrange(quantized, 'b 1 d -> b d')
+            indices = rearrange(indices, 'b 1 ... -> b ...')
+
+        reconstructed = self.decoder(quantized)
+        return Return(reconstructed, indices, aux_loss)
+
+    @torch.no_grad()
+    def encode(self, x, return_signs = False):
+        with temp_eval(self):
+            latents = self.encoder(x)
+            is_2d = latents.ndim == 2
+
+            if is_2d:
+                latents = rearrange(latents, 'b d -> b 1 d')
+
+            quantized, indices, _ = self.lfq(latents)
+
+            if is_2d:
+                quantized = rearrange(quantized, 'b 1 d -> b d')
+
+            if return_signs:
+                return torch.where(quantized > 0, 1.0, -1.0)
+
+            return (quantized > 0).float()
+
+    @torch.no_grad()
+    def decode_bits(self, bits):
+        """
+        Converts binary bits (0/1 or -1/+1) to codebook indices and uses lfq.indices_to_codes
+        to decode back to data space.
+        """
+        with temp_eval(self):
+            bits = (bits > 0).int()
+            bits, inverse = pack_with_inverse(bits, '* d')
+
+            codebook_dim = self.lfq.codebook_dim
+            num_codebooks = self.lfq.num_codebooks
+
+            if num_codebooks > 1 and bits.shape[-1] == codebook_dim * num_codebooks:
+                bits = rearrange(bits, 'b (c d) -> b c d', c = num_codebooks, d = codebook_dim)
+
+            if bits.ndim == 2:
+                bits = rearrange(bits, 'b d -> b 1 d')
+
+            mask = 2 ** torch.arange(codebook_dim - 1, -1, -1, device = self.device)
+            indices = reduce(bits * mask, '... d -> ...', 'sum')
+
+            if not self.lfq.keep_num_codebooks_dim and indices.ndim >= 2 and indices.shape[-1] == 1:
+                indices = rearrange(indices, '... 1 -> ...')
+
+            codes = self.lfq.indices_to_codes(indices)
+
+            if codes.ndim == 3 and codes.shape[1] == 1:
+                codes = rearrange(codes, 'b 1 d -> b d')
+
+            codes = inverse(codes, '* d')
+            return self.decoder(codes)
+
+    # genetic algorithm helpers
+
+    def init_random_population(self, pop_size = None, shape = None, device = None, is_sign = False):
+        pop_size = default(pop_size, self.pop_size)
+        device = default(device, self.device)
+        bits = torch.randint(0, 2, (pop_size, *shape), device = device).float()
+
+        if is_sign:
+            return torch.where(bits > 0, 1.0, -1.0)
+
+        return bits
+
+    def init_population_from_data(self, x, pop_size = None, mutation_rate = None, is_sign = False):
+        pop_size = default(pop_size, self.pop_size)
+        mutation_rate = default(mutation_rate, self.mutation_rate)
+        bits = self.encode(x, return_signs = is_sign)
+
+        num_samples, *_ = bits.shape
+
+        if num_samples < pop_size:
+            repeats = (pop_size + num_samples - 1) // num_samples
+            bits = repeat(bits, 'b ... -> (r b) ...', r = repeats)[:pop_size]
+        else:
+            bits = bits[:pop_size]
+
+        if mutation_rate > 0:
+            bits = self.mutate(bits, mutation_rate = mutation_rate, is_sign = is_sign)
+
+        return bits
+
+    def uniform_crossover(self, parent1, parent2):
+        crossover_mask = torch.rand_like(parent1.float()) < 0.5
+        return torch.where(crossover_mask, parent1, parent2)
+
+    def mutate(self, population, mutation_rate = None, is_sign = False):
+        mutation_rate = default(mutation_rate, self.mutation_rate)
+        flip_mask = torch.rand_like(population.float()) < mutation_rate
+
+        if is_sign:
+            return torch.where(flip_mask, -population, population)
+
+        return torch.where(flip_mask, 1.0 - population, population)
+
+    def tournament_selection(self, population, fitnesses, tournament_size = None):
+        tournament_size = default(tournament_size, self.tournament_size)
+        pop_size, *_ = population.shape
+
+        contenders = torch.randint(0, pop_size, (pop_size, tournament_size), device = self.device)
+        winner_indices = contenders.gather(1, fitnesses[contenders].argmax(dim = -1, keepdim = True))
+        return population[rearrange(winner_indices, 'b 1 -> b')]
+
+    @torch.no_grad()
+    def step(
+        self,
+        pop_bits,
+        fitness_fn,
+        tournament_size = None,
+        mutation_rate = None,
+        elitism_count = None,
+        is_sign = False,
+        batch_size = None,
+        **kwargs
+    ):
+        tournament_size = default(tournament_size, self.tournament_size)
+        mutation_rate = default(mutation_rate, self.mutation_rate)
+        elitism_count = default(elitism_count, self.elitism_count)
+
+        pop_size, *_ = pop_bits.shape
+        num_offspring = pop_size - elitism_count
+
+        # evaluate fitness
+
+        if exists(batch_size):
+            decoded_chunks = [self.decode_bits(chunk) for chunk in pop_bits.split(batch_size)]
+            decoded = torch.cat(decoded_chunks, dim = 0)
+        else:
+            decoded = self.decode_bits(pop_bits)
+
+        fitnesses = fitness_fn(decoded, pop_bits)
+        if not is_tensor(fitnesses):
+            fitnesses = torch.tensor(fitnesses, device = self.device, dtype = torch.float32)
+
+        # sort by fitness descending and preserve elites
+
+        sorted_indices = torch.argsort(fitnesses, descending = True)
+        elites = pop_bits[sorted_indices[:elitism_count]].clone()
+
+        # tournament selection: 2 parent tournaments for each of the num_offspring children needed
+
+        contenders = torch.randint(0, pop_size, (num_offspring, 2, tournament_size), device = self.device)
+        winner_rel_idx = fitnesses[contenders].argmax(dim = -1, keepdim = True)
+        parent_indices = rearrange(contenders.gather(-1, winner_rel_idx), 'n p 1 -> n p')
+
+        # crossover: 1 child produced per 2 parents
+
+        p1, p2 = rearrange(pop_bits[parent_indices], 'n p ... -> p n ...')
+        offspring = self.uniform_crossover(p1, p2)
+
+        # mutation all at once
+
+        offspring = self.mutate(offspring, mutation_rate = mutation_rate, is_sign = is_sign)
+
+        next_pop = torch.cat([elites, offspring], dim = 0)
+        return next_pop, fitnesses
+
+    @torch.no_grad()
+    def evolve(
+        self,
+        fitness_fn,
+        pop_bits = None,
+        pop_size = None,
+        shape = None,
+        generations = None,
+        is_sign = False,
+        return_best_decoded = False,
+        **step_kwargs
+    ):
+        pop_size = default(pop_size, self.pop_size)
+        generations = default(generations, self.generations)
+
+        if not exists(pop_bits):
+            assert exists(shape), 'shape must be provided if pop_bits is not supplied to evolve()'
+            pop_bits = self.init_random_population(pop_size, shape, is_sign = is_sign)
+
+        best_fitness = float('-inf')
+        best_gene = None
+
+        for _ in range(generations):
+            pop_bits, fitnesses = self.step(
+                pop_bits,
+                fitness_fn,
+                is_sign = is_sign,
+                **step_kwargs
+            )
+
+            max_fit_idx = torch.argmax(fitnesses)
+            max_fit = fitnesses[max_fit_idx].item()
+
+            if max_fit > best_fitness:
+                best_fitness = max_fit
+                best_gene = pop_bits[max_fit_idx].clone()
+
+            best_decoded = None
+            if return_best_decoded:
+                best_decoded = self.decode_bits(rearrange(best_gene, '... -> 1 ...'))
+                best_decoded = rearrange(best_decoded, '1 ... -> ...')
+
+            yield Result(pop_bits, best_gene, best_fitness, best_decoded)

{vector_quantize_pytorch-1.31.0 → vector_quantize_pytorch-1.31.2}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -654,14 +654,15 @@ class Codebook(Module):
 
         if needs_codebook_dim:
             x = rearrange(x, '... -> 1 ...')
+            embed_ind = rearrange(embed_ind, '... -> 1 ...')
 
         dtype = x.dtype
-        flatten, unpack_one = pack_one(x, 'h * d')
+        flatten, _ = pack_one(x, 'h * d')
 
         if exists(mask):
             mask = repeat(mask, 'b n -> c (b h n)', c = flatten.shape[0], h = flatten.shape[-2] // (mask.shape[0] * mask.shape[1]))
 
-        embed_ind, _ = pack([embed_ind], 'h *')
+        embed_ind, _ = pack_one(embed_ind, 'h *')
         embed_ind = embed_ind.masked_fill(embed_ind == -1, 0)
         embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
 
@@ -697,6 +698,8 @@ class Codebook(Module):
         dtype = x.dtype
         flatten, unpack_one = pack_one(x, 'h * d')
 
+        is_topk = exists(topk)
+
         if exists(mask):
             mask = repeat(mask, 'b n -> c (b h n)', c = flatten.shape[0], h = flatten.shape[-2] // (mask.shape[0] * mask.shape[1]))
 
@@ -746,19 +749,13 @@ class Codebook(Module):
 
         embed_ind, embed_onehot = self.gumbel_sample(dist, dim = -1, topk = topk, temperature = sample_codebook_temp, training = self.training)
 
-        if exists(topk):
-            embed_ind = unpack_one(embed_ind, 'h * k')
-        else:
-            embed_ind = unpack_one(embed_ind, 'h *')
+        embed_ind = unpack_one(embed_ind, 'h * k' if is_topk else 'h *')
 
         if exists(codebook_transform_fn):
             transformed_embed = unpack_one(transformed_embed, 'h * c d')
 
-        if self.training:
-            if exists(topk):
-                unpacked_onehot = unpack_one(embed_onehot, 'h * k c')
-            else:
-                unpacked_onehot = unpack_one(embed_onehot, 'h * c')
+        if self.training or is_topk:
+            unpacked_onehot = unpack_one(embed_onehot, 'h * k c' if is_topk else 'h * c')
 
             if exists(codebook_transform_fn):
                 quantize = einsum('h b n ... c, h b n c d -> h b n ... d', unpacked_onehot, transformed_embed)
@@ -780,7 +777,7 @@ class Codebook(Module):
                 repeated_embed_ind = repeat(embed_ind, 'h b n -> h b n d', d = embed.shape[-1])
                 quantize = repeated_embed.gather(-2, repeated_embed_ind)
 
-        if self.training and update_usage and not freeze_codebook and not exists(topk):
+        if self.training and update_usage and not freeze_codebook and not is_topk:
             self.update_codebook(flatten, embed_onehot, mask = mask, ema_update_weight = ema_update_weight, accum_ema_update = accum_ema_update, ema_update = ema_update)
 
         if needs_codebook_dim:
@@ -1107,6 +1104,18 @@ class VectorQuantize(Module):
     ):
         orig_input, input_requires_grad = x, x.requires_grad
 
+        is_topk = exists(topk)
+
+        if is_topk:
+            topk = min(int(topk), self.codebook_size)
+
+            assert (
+                self.heads == 1
+                and topk > 0
+                and not (self.commitment_use_cross_entropy_loss and self.training)
+                and not (exists(self.in_place_codebook_optimizer) and self.training)
+            ), 'topk codes only supported for single-head codebooks without cross entropy commitment loss or in-place codebook optimizer'
+
         # freezing codebook
 
         freeze_codebook = default(freeze_codebook, self.freeze_codebook)
@@ -1128,7 +1137,7 @@ class VectorQuantize(Module):
 
         shape, dtype, device, heads, is_multiheaded, codebook_size, return_loss = x.shape, x.dtype, x.device, self.heads, self.heads > 1, self.codebook_size, exists(indices)
 
-        need_transpose = not self.channel_last and not self.accept_image_fmap and not self.accept_3d_fmap
+        need_transpose = not self.channel_last and not self.accept_image_fmap and not self.accept_3d_fmap and not only_one
         should_inplace_optimize = exists(self.in_place_codebook_optimizer)
 
         # rearrange inputs
@@ -1167,7 +1176,7 @@ class VectorQuantize(Module):
             codebook_transform_fn = codebook_transform_fn,
             ema_update_weight = ema_update_weight,
             accum_ema_update = accum_ema_update,
-            ema_update = ema_update and not exists(topk),
+            ema_update = ema_update and not is_topk,
             topk = topk
         )
 
@@ -1209,17 +1218,14 @@ class VectorQuantize(Module):
             codebook_forward_kwargs.update(update_usage = False)
             quantize, embed_ind, distances = self._codebook(x, **codebook_forward_kwargs)
 
+        if is_topk:
+            x = repeat(x, '... d -> ... k d', k = topk)
+
         if self.training:
-            # determine code to use for commitment loss
             maybe_detach = torch.detach if not self.learnable_codebook or freeze_codebook else identity
 
             commit_quantize = maybe_detach(quantize)
 
-            # maybe expand input if returning topk codes
-
-            if exists(topk):
-                x = repeat(x, '... d -> ... k d', k = topk)
-
             # spare rotation trick calculation if inputs do not need gradients
 
             if input_requires_grad and self.route_gradients_to_input:
@@ -1279,7 +1285,19 @@ class VectorQuantize(Module):
 
         # aggregate loss
 
-        loss = tensor(0., device = device, requires_grad = self.training)
+        loss = self.zero
+
+        if self.training:
+            loss.requires_grad_()
+
+        if is_topk and not self.training:
+            commit_loss = F.mse_loss(quantize, x, reduction = 'none')
+            commit_loss = reduce(commit_loss, '... k d -> ... k', 'mean')
+
+            if exists(mask):
+                commit_loss = einx.where('..., ... k, -> ... k', mask, commit_loss, 0.)
+
+            loss = commit_loss * self.commitment_weight if self.has_commitment_loss else commit_loss
 
         if self.training:
             # calculate codebook diversity loss (negative of entropy) if needed
@@ -1304,11 +1322,8 @@ class VectorQuantize(Module):
 
                     commit_loss = calculate_ce_loss(embed_ind)
                 else:
-                    if exists(topk):
-                        # handle special case when returning topk
-
-                        repeated_input = repeat(orig_input, '... d -> ... k d', k = topk)
-                        commit_loss = F.mse_loss(commit_quantize, repeated_input, reduction = 'none')
+                    if is_topk:
+                        commit_loss = F.mse_loss(commit_quantize, x, reduction = 'none')
                         commit_loss = reduce(commit_loss, '... k d -> ... k', 'mean')
 
                         if exists(mask):
@@ -1347,6 +1362,18 @@ class VectorQuantize(Module):
                 orthogonal_reg_loss = orthogonal_loss_fn(codebook)
                 loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight
 
+        # reshape the topk loss to have the same shape as the embedding indices, if it is a per token loss
+
+        if is_topk and loss.ndim > 1:
+            if only_one:
+                loss = rearrange(loss, 'b 1 ... -> b ...')
+
+            if self.accept_image_fmap:
+                loss = rearrange(loss, 'b (h w) ... -> b h w ...', h = height, w = width)
+
+            if self.accept_3d_fmap:
+                loss = rearrange(loss, "b (d h w) ... -> b d h w ...", d = depth, h = height, w = width)
+
         # handle multi-headed quantized embeddings
 
         if is_multiheaded:
@@ -1362,16 +1389,16 @@ class VectorQuantize(Module):
         # rearrange quantized embeddings
 
         if need_transpose:
-            quantize = rearrange(quantize, 'b n d -> b d n')
+            quantize = rearrange(quantize, 'b n ... d -> b d n ...')
 
         if self.accept_image_fmap:
-            quantize = rearrange(quantize, 'b (h w) c -> b c h w', h = height, w = width)
+            quantize = rearrange(quantize, 'b (h w) ... c -> b c h w ...', h = height, w = width)
 
         if self.accept_3d_fmap:
-            quantize = rearrange(quantize, "b (d h w) c -> b c d h w", d=depth, h=height, w=width)
+            quantize = rearrange(quantize, "b (d h w) ... c -> b c d h w ...", d=depth, h=height, w=width)
 
         if only_one:
-            quantize = rearrange(quantize, 'b 1 d -> b d')
+            quantize = rearrange(quantize, 'b 1 ... d -> b ... d')
 
         # if masking, only return quantized for where mask has True