egogym 0.1.0__py3-none-any.whl

baselines/rum/models/bet/vqvae/residual_vq.py

@@ -0,0 +1,346 @@
+from math import ceil
+from functools import partial
+from itertools import zip_longest
+from random import randrange
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from baselines.rum.models.bet.vqvae.vector_quantize_pytorch import VectorQuantize
+
+from einops import rearrange, repeat, pack, unpack
+
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def round_up_multiple(num, mult):
+    return ceil(num / mult) * mult
+
+
+# main class
+
+
+class ResidualVQ(nn.Module):
+    """Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf"""
+
+    def __init__(
+        self,
+        *,
+        dim,
+        num_quantizers,
+        codebook_dim=None,
+        shared_codebook=False,
+        heads=1,
+        quantize_dropout=False,
+        quantize_dropout_cutoff_index=0,
+        quantize_dropout_multiple_of=1,
+        accept_image_fmap=False,
+        eval=False,
+        **kwargs,
+    ):
+        super().__init__()
+        assert heads == 1, "residual vq is not compatible with multi-headed codes"
+        codebook_dim = default(codebook_dim, dim)
+        codebook_input_dim = codebook_dim * heads
+
+        requires_projection = codebook_input_dim != dim
+        self.project_in = (
+            nn.Linear(dim, codebook_input_dim) if requires_projection else nn.Identity()
+        )
+        self.project_out = (
+            nn.Linear(codebook_input_dim, dim) if requires_projection else nn.Identity()
+        )
+
+        self.num_quantizers = num_quantizers
+
+        self.accept_image_fmap = accept_image_fmap
+        self.layers = nn.ModuleList(
+            [
+                VectorQuantize(
+                    dim=codebook_dim,
+                    codebook_dim=codebook_dim,
+                    accept_image_fmap=accept_image_fmap,
+                    eval=eval,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+        if eval:
+            self.layers.eval()
+
+        self.quantize_dropout = quantize_dropout and num_quantizers > 1
+
+        assert quantize_dropout_cutoff_index >= 0
+
+        self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index
+        self.quantize_dropout_multiple_of = quantize_dropout_multiple_of  # encodec paper proposes structured dropout, believe this was set to 4
+
+        if not shared_codebook:
+            return
+
+        first_vq, *rest_vq = self.layers
+        codebook = first_vq._codebook
+
+        for vq in rest_vq:
+            vq._codebook = codebook
+
+    @property
+    def codebooks(self):
+        codebooks = [layer._codebook.embed for layer in self.layers]
+        codebooks = torch.stack(codebooks, dim=0)
+        codebooks = rearrange(codebooks, "q 1 c d -> q c d")
+        return codebooks
+
+    def get_codes_from_indices(self, indices):
+        batch, quantize_dim = indices.shape[0], indices.shape[-1]
+
+        # may also receive indices in the shape of 'b h w q' (accept_image_fmap)
+
+        indices, ps = pack([indices], "b * q")
+
+        # because of quantize dropout, one can pass in indices that are coarse
+        # and the network should be able to reconstruct
+
+        if quantize_dim < self.num_quantizers:
+            assert self.quantize_dropout > 0.0, (
+                "quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations"
+            )
+            indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value=-1)
+
+        # get ready for gathering
+
+        codebooks = repeat(self.codebooks, "q c d -> q b c d", b=batch)
+        gather_indices = repeat(indices, "b n q -> q b n d", d=codebooks.shape[-1])
+
+        # take care of quantizer dropout
+
+        mask = gather_indices == -1.0
+        gather_indices = gather_indices.masked_fill(
+            mask, 0
+        )  # have it fetch a dummy code to be masked out later
+
+        all_codes = codebooks.gather(2, gather_indices)  # gather all codes
+
+        # mask out any codes that were dropout-ed
+
+        all_codes = all_codes.masked_fill(mask, 0.0)
+
+        # if (accept_image_fmap = True) then return shape (quantize, batch, height, width, dimension)
+
+        (all_codes,) = unpack(all_codes, ps, "q b * d")
+
+        return all_codes
+
+    def draw_logits_forward(self, encoding_logits):
+        # encoding_indices : dim1 = batch_size  dim2 = 4 (number of groups) dim3 = vq dict size (header)
+        encoding_logits = encoding_logits.to(self.device)
+        bs = encoding_logits.shape[0]
+        quantized = torch.zeros((bs, self.codebooks.shape[-1])).to(self.device)
+        for q in range(encoding_logits.shape[1]):
+            quantized += torch.matmul(encoding_logits[:, q], self.codebooks[q]).to(
+                self.device
+            )
+        return quantized
+
+    def forward(
+        self, x, indices=None, return_all_codes=False, sample_codebook_temp=None
+    ):
+        num_quant, quant_dropout_multiple_of, return_loss, device = (
+            self.num_quantizers,
+            self.quantize_dropout_multiple_of,
+            exists(indices),
+            x.device,
+        )
+
+        x = self.project_in(x)
+
+        assert not (self.accept_image_fmap and exists(indices))
+
+        quantized_out = 0.0
+        residual = x
+
+        all_losses = []
+        all_indices = []
+
+        if return_loss:
+            assert not torch.any(indices == -1), (
+                "some of the residual vq indices were dropped out. please use indices derived when the module is in eval mode to derive cross entropy loss"
+            )
+            ce_losses = []
+
+        should_quantize_dropout = (
+            self.training and self.quantize_dropout and not return_loss
+        )
+
+        # sample a layer index at which to dropout further residual quantization
+        # also prepare null indices and loss
+
+        if should_quantize_dropout:
+            rand_quantize_dropout_index = randrange(
+                self.quantize_dropout_cutoff_index, num_quant
+            )
+
+            if quant_dropout_multiple_of != 1:
+                rand_quantize_dropout_index = (
+                    round_up_multiple(
+                        rand_quantize_dropout_index + 1, quant_dropout_multiple_of
+                    )
+                    - 1
+                )
+
+            null_indices_shape = (
+                (x.shape[0], *x.shape[-2:])
+                if self.accept_image_fmap
+                else tuple(x.shape[:2])
+            )
+            null_indices = torch.full(
+                null_indices_shape, -1.0, device=device, dtype=torch.long
+            )
+            null_loss = torch.full((1,), 0.0, device=device, dtype=x.dtype)
+
+        # go through the layers
+
+        for quantizer_index, layer in enumerate(self.layers):
+            if (
+                should_quantize_dropout
+                and quantizer_index > rand_quantize_dropout_index
+            ):
+                all_indices.append(null_indices)
+                all_losses.append(null_loss)
+                continue
+
+            layer_indices = None
+            if return_loss:
+                layer_indices = indices[..., quantizer_index]
+
+            quantized, *rest = layer(
+                residual,
+                indices=layer_indices,
+                sample_codebook_temp=sample_codebook_temp,
+            )
+
+            residual = residual - quantized.detach()
+            quantized_out = quantized_out + quantized
+
+            if return_loss:
+                ce_loss = rest[0]
+                ce_losses.append(ce_loss)
+                continue
+
+            embed_indices, loss = rest
+
+            all_indices.append(embed_indices)
+            all_losses.append(loss)
+
+        # project out, if needed
+
+        quantized_out = self.project_out(quantized_out)
+
+        # whether to early return the cross entropy loss
+
+        if return_loss:
+            return quantized_out, sum(ce_losses)
+
+        # stack all losses and indices
+
+        all_losses, all_indices = map(
+            partial(torch.stack, dim=-1), (all_losses, all_indices)
+        )
+
+        ret = (quantized_out, all_indices, all_losses)
+
+        if return_all_codes:
+            # whether to return all codes from all codebooks across layers
+            all_codes = self.get_codes_from_indices(all_indices)
+
+            # will return all codes in shape (quantizer, batch, sequence length, codebook dimension)
+            ret = (*ret, all_codes)
+
+        return ret
+
+
+# grouped residual vq
+
+
+class GroupedResidualVQ(nn.Module):
+    def __init__(self, *, dim, groups=1, accept_image_fmap=False, **kwargs):
+        super().__init__()
+        self.dim = dim
+        self.groups = groups
+        assert (dim % groups) == 0
+        dim_per_group = dim // groups
+
+        self.accept_image_fmap = accept_image_fmap
+
+        self.rvqs = nn.ModuleList([])
+
+        for _ in range(groups):
+            self.rvqs.append(
+                ResidualVQ(
+                    dim=dim_per_group, accept_image_fmap=accept_image_fmap, **kwargs
+                )
+            )
+
+    @property
+    def codebooks(self):
+        return torch.stack(tuple(rvq.codebooks for rvq in self.rvqs))
+
+    def get_codes_from_indices(self, indices):
+        codes = tuple(
+            rvq.get_codes_from_indices(chunk_indices)
+            for rvq, chunk_indices in zip(self.rvqs, indices)
+        )
+        return torch.stack(codes)
+
+    def forward(
+        self, x, indices=None, return_all_codes=False, sample_codebook_temp=None
+    ):
+        shape = x.shape
+        split_dim = 1 if self.accept_image_fmap else -1
+        assert shape[split_dim] == self.dim
+
+        # split the feature dimension into groups
+
+        x = x.chunk(self.groups, dim=split_dim)
+
+        indices = default(indices, tuple())
+        return_ce_loss = len(indices) > 0
+        assert len(indices) == 0 or len(indices) == self.groups
+
+        forward_kwargs = dict(
+            return_all_codes=return_all_codes, sample_codebook_temp=sample_codebook_temp
+        )
+
+        # invoke residual vq on each group
+
+        out = tuple(
+            rvq(chunk, indices=chunk_indices, **forward_kwargs)
+            for rvq, chunk, chunk_indices in zip_longest(self.rvqs, x, indices)
+        )
+        out = tuple(zip(*out))
+
+        # if returning cross entropy loss to rvq codebooks
+
+        if return_ce_loss:
+            quantized, ce_losses = out
+            return torch.cat(quantized, dim=split_dim), sum(ce_losses)
+
+        # otherwise, get all the zipped outputs and combine them
+
+        quantized, all_indices, commit_losses, *maybe_all_codes = out
+
+        quantized = torch.cat(quantized, dim=split_dim)
+        all_indices = torch.stack(all_indices)
+        commit_losses = torch.stack(commit_losses)
+
+        ret = (quantized, all_indices, commit_losses, *maybe_all_codes)
+        return ret