Method forward

easyvolcap/utils/vq_utils.py:291–333

(self, x, weight=None, verbose=False)

    @autocast(enabled=False)
    def forward(self, x, weight=None, verbose=False):
        if weight is not None:
            weight = weight * weight.numel() / weight.sum()
        needs_codebook_dim = x.ndim < 4

        x = x.float()

        if needs_codebook_dim:
            x = rearrange(x, '... -> 1 ...')

        shape, dtype = x.shape, x.dtype
        flatten = rearrange(x, 'h ... d -> h (...) d')
        self.init_embed_(flatten)
        embed = self.embed if not self.learnable_codebook else self.embed.detach()
        dist = -torch.cdist(flatten, embed, p=2)
        embed_ind = gumbel_sample(dist, dim=-1, temperature=self.sample_codebook_temp)
        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
        embed_ind = embed_ind.view(*shape[:-1])
        quantize = batched_embedding(embed_ind, self.embed)

        if self.training:

            if weight is not None:
                cluster_size = (embed_onehot * weight).sum(dim=1)
            else:
                cluster_size = embed_onehot.sum(dim=1)
            self.all_reduce_fn(cluster_size)
            ema_inplace(self.cluster_size, cluster_size, self.decay)

            if weight is not None:
                embed_sum = einsum('h n d, h n c -> h c d', flatten * weight, embed_onehot)
            else:
                embed_sum = einsum('h n d, h n c -> h c d', flatten, embed_onehot)
            self.all_reduce_fn(embed_sum)
            cluster_size = laplace_smoothing(self.cluster_size, self.codebook_size, self.eps) * self.cluster_size.sum()

            ema_inplace(self.embed, embed_sum / rearrange(cluster_size, '... -> ... 1'), self.decay)
            self.expire_codes_(x, verbose)

        if needs_codebook_dim:
            quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))

        return quantize, embed_ind

# main class