Publish pinned visual components with exact source payloads and native numerical validation

b9a3411 verified 3 days ago

20.5 kB

# --------------------------------------------------------
# BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers (https://arxiv.org/abs/2208.06366)
# Github source: https://github.com/microsoft/unilm/tree/master/beitv2
# Copyright (c) 2022 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# By Zhiliang Peng
# Based on VQGAN code bases
# https://github.com/CompVis/taming-transformers
# --------------------------------------------------------'

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as distributed
from einops import rearrange, repeat


def l2norm(t):
    return F.normalize(t, p = 2, dim = -1)

def ema_inplace(moving_avg, new, decay):
    moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))

def sample_vectors(samples, num):
    num_samples, device = samples.shape[0], samples.device

    if num_samples >= num:
        indices = torch.randperm(num_samples, device = device)[:num]
    else:
        indices = torch.randint(0, num_samples, (num,), device = device)

    return samples[indices]

def kmeans(samples, num_clusters, num_iters = 10, use_cosine_sim = False):
    
    # rng_state = torch.get_rng_state()
    # torch.manual_seed(42)
    
    dim, dtype, device = samples.shape[-1], samples.dtype, samples.device

    means = sample_vectors(samples, num_clusters)

    for i in range(num_iters):
        if use_cosine_sim:
            dists = samples @ means.t()
        else:
            diffs = rearrange(samples, 'n d -> n () d') \
                    - rearrange(means, 'c d -> () c d')
            dists = -(diffs ** 2).sum(dim = -1)

        buckets = dists.max(dim = -1).indices
        bins = torch.bincount(buckets, minlength = num_clusters)
        zero_mask = bins == 0
        bins_min_clamped = bins.masked_fill(zero_mask, 1)

        new_means = buckets.new_zeros(num_clusters, dim, dtype = dtype)
        new_means.scatter_add_(0, repeat(buckets, 'n -> n d', d = dim), samples)
        new_means = new_means / bins_min_clamped[..., None]

        if use_cosine_sim:
            new_means = l2norm(new_means)

        means = torch.where(zero_mask[..., None], means, new_means)
        print(f"Iteration {i + 1}/{num_iters} - Zero clusters: {zero_mask.sum().item()}")
    
    # torch.set_rng_state(rng_state)
    return means, bins


class EmbeddingEMA(nn.Module):
    def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5, kmeans_init=True, codebook_init_path=''):
        super().__init__()
        self.num_tokens = num_tokens
        self.codebook_dim = codebook_dim
        self.decay = decay
        self.eps = eps 
        if codebook_init_path == '':   
            if not kmeans_init:
                weight = torch.randn(num_tokens, codebook_dim)
                weight = l2norm(weight)
            else:
                weight = torch.zeros(num_tokens, codebook_dim)
            self.register_buffer('initted', torch.Tensor([not kmeans_init]))
        else:
            print(f"load init codebook weight from {codebook_init_path}")
            codebook_ckpt_weight = torch.load(codebook_init_path, map_location='cpu')
            weight = codebook_ckpt_weight.clone()
            self.register_buffer('initted', torch.Tensor([True]))
            
        self.weight = nn.Parameter(weight, requires_grad = False)
        self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad = False)
        self.embed_avg = nn.Parameter(weight.clone(), requires_grad = False)
        self.update = True

    @torch.jit.ignore
    def init_embed_(self, data):
        if self.initted:
            return
        print("Performing Kemans init for codebook")
        embed, cluster_size = kmeans(data, self.num_tokens, 10, use_cosine_sim = True)
        self.weight.data.copy_(embed)
        self.cluster_size.data.copy_(cluster_size)
        self.initted.data.copy_(torch.Tensor([True]))

    @torch.jit.ignore
    def init_embed_split(self, data, split):
        if self.initted:
            return
        print("Performing Kemans init for codebook")
        embed1, cluster_size1 = kmeans(data[:, :split[0]], self.num_tokens, 10, use_cosine_sim = True)
        embed2, cluster_size2 = kmeans(data[:, split[0]:], self.num_tokens, 10, use_cosine_sim = True)
        embed = torch.cat([embed1, embed2], dim=-1)
        cluster_size = (cluster_size1 + cluster_size2) / 2.
        self.weight.data.copy_(embed)
        self.cluster_size.data.copy_(cluster_size)
        self.initted.data.copy_(torch.Tensor([True]))
    
    @torch.jit.ignore
    def init_embed_with_ind(self, data, inds):
        if self.initted:
            return
        print("Performing Kemans init for codebook")
        embed, cluster_size = kmeans(data, self.num_tokens, 10, use_cosine_sim = True)
        self.weight.data.copy_(embed)
        self.cluster_size.data.copy_(cluster_size)
        self.initted.data.copy_(torch.Tensor([True]))
    def forward(self, embed_id):
        return F.embedding(embed_id, self.weight)

    def cluster_size_ema_update(self, new_cluster_size):
        self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)

    def embed_avg_ema_update(self, new_embed_avg): 
        self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)

    def weight_update(self, num_tokens):
        n = self.cluster_size.sum()
        smoothed_cluster_size = (
                (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
            )
        #normalize embedding average with smoothed cluster size
        embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
        # embed_normalized = l2norm(self.embed_avg / smoothed_cluster_size.unsqueeze(1))
        self.weight.data.copy_(embed_normalized)   

def norm_ema_inplace(moving_avg, new, decay):
    moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))
    moving_avg.data.copy_(l2norm(moving_avg.data))

class NormEMAVectorQuantizer(nn.Module):
    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5, 
                statistic_code_usage=True, kmeans_init=False, codebook_init_path=''):
        super().__init__()
        self.codebook_dim = embedding_dim
        self.num_tokens = n_embed
        self.beta = beta
        self.decay = decay
        
        # learnable = True if orthogonal_reg_weight > 0 else False
        self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps, kmeans_init, codebook_init_path)
        
        self.statistic_code_usage = statistic_code_usage
        if statistic_code_usage:
            self.register_buffer('cluster_size', torch.zeros(n_embed))
        if distributed.is_available() and distributed.is_initialized():
            print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
            self.all_reduce_fn = distributed.all_reduce
        else:
            self.all_reduce_fn = nn.Identity()
    
    def reset_cluster_size(self, device):
        if self.statistic_code_usage:
            self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
            self.cluster_size = self.cluster_size.to(device)

    def forward(self, z):
        # reshape z -> (batch, height, width, channel) and flatten
        #z, 'b c h w -> b h w c'
        z = rearrange(z, 'b c h w -> b h w c')
        z = l2norm(z)
        z_flattened = z.reshape(-1, self.codebook_dim)
        
        self.embedding.init_embed_(z_flattened)
        
        d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \
            self.embedding.weight.pow(2).sum(dim=1) - 2 * \
            torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n'
        
        encoding_indices = torch.argmin(d, dim=1)

        z_q = self.embedding(encoding_indices).view(z.shape)
        
        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)     
        
        if not self.training:
            with torch.no_grad():
                cluster_size = encodings.sum(0)
                self.all_reduce_fn(cluster_size)
                ema_inplace(self.cluster_size, cluster_size, self.decay)
        
        if self.training and self.embedding.update:
            #EMA cluster size

            bins = encodings.sum(0)
            self.all_reduce_fn(bins)

            ema_inplace(self.cluster_size, bins, self.decay)

            zero_mask = (bins == 0)
            bins = bins.masked_fill(zero_mask, 1.)

            embed_sum = z_flattened.t() @ encodings
            self.all_reduce_fn(embed_sum)
                        
            embed_normalized = (embed_sum / bins.unsqueeze(0)).t()
            embed_normalized = l2norm(embed_normalized)
            
            embed_normalized = torch.where(zero_mask[..., None], self.embedding.weight,
                                           embed_normalized)
            norm_ema_inplace(self.embedding.weight, embed_normalized, self.decay)

        loss = self.beta * F.mse_loss(z_q.detach(), z) 
        
        z_q = z + (z_q - z).detach()

        #z_q, 'b h w c -> b c h w'
        z_q = rearrange(z_q, 'b h w c -> b c h w')
        return z_q, loss, encoding_indices
    


class VectorQuantizer(nn.Module):
    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5, 
                statistic_code_usage=True, kmeans_init=False, codebook_init_path=''):
        super().__init__()
        self.codebook_dim = embedding_dim
        self.num_tokens = n_embed
        self.beta = beta
        decay = 0.
        self.decay = decay
        
        # learnable = True if orthogonal_reg_weight > 0 else False
        self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps, kmeans_init, codebook_init_path)
        
        self.statistic_code_usage = statistic_code_usage
        if statistic_code_usage:
            self.register_buffer('cluster_size', torch.zeros(n_embed))
        if distributed.is_available() and distributed.is_initialized():
            print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
            self.all_reduce_fn = distributed.all_reduce
        else:
            self.all_reduce_fn = nn.Identity()
    
    def reset_cluster_size(self, device):
        if self.statistic_code_usage:
            self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
            self.cluster_size = self.cluster_size.to(device)

    def forward(self, bhw, encoding_indices, z=None):
        embed_size = [bhw[0], bhw[1], bhw[2], -1]
        z_q = self.embedding(encoding_indices).view(embed_size)
        
        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z_q.dtype)   
        
        if not self.training:
            with torch.no_grad():
                cluster_size = encodings.sum(0)
                self.all_reduce_fn(cluster_size)
                ema_inplace(self.cluster_size, cluster_size, self.decay)
        if z is not None:
            #EMA cluster size
            z = rearrange(z, 'b c h w -> b h w c')
            loss = self.beta * F.mse_loss(z_q, z.detach()) 

        else:
            loss = torch.tensor([0.]).mean().to(z_q)
        z_q = rearrange(z_q, 'b h w c -> b c h w')
        return z_q, loss, encoding_indices
    


class CVectorQuantiser(nn.Module):
    """
    Improved version over vector quantiser, with the dynamic initialisation
    for these unoptimised "dead" points.
    num_embed: number of codebook entry
    embed_dim: dimensionality of codebook entry
    beta: weight for the commitment loss
    distance: distance for looking up the closest code
    anchor: anchor sampled methods
    first_batch: if true, the offline version of our model
    contras_loss: if true, use the contras_loss to further improve the performance
    """
    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
                 statistic_code_usage=True,  kmeans_init=True, codebook_init_path='', 
                 distance='l2', anchor='closest', first_batch=False, contras_loss=False):
        super().__init__()

        self.num_embed = n_embed
        self.n_e = n_embed
        self.num_tokens = n_embed
        self.embed_dim = embedding_dim
        
        self.beta = beta
        self.distance = distance
        self.anchor = anchor
        self.first_batch = first_batch
        self.contras_loss = contras_loss
        self.decay = decay
        self.init = False

        self.pool = FeaturePool(self.num_embed, self.embed_dim)
        # self.embedding = nn.Embedding(self.num_embed, self.embed_dim)
        # self.embedding.weight.data.uniform_(-1.0 / self.num_embed, 1.0 / self.num_embed)
        self.embedding = EmbeddingEMA(self.num_tokens, self.embed_dim, decay, eps, kmeans_init, codebook_init_path)
        
        self.register_buffer("embed_prob", torch.zeros(self.num_embed))
        self.statistic_code_usage = statistic_code_usage
        if statistic_code_usage:
            self.register_buffer('cluster_size', torch.zeros(n_embed))
        if distributed.is_available() and distributed.is_initialized():
            print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
            self.all_reduce_fn = distributed.all_reduce
        else:
            self.all_reduce_fn = nn.Identity()

    def reset_cluster_size(self, device):
        if self.statistic_code_usage:
            self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
            self.cluster_size = self.cluster_size.to(device)

    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
        # assert temp is None or temp==1.0, "Only for interface compatible with Gumbel"
        # assert rescale_logits==False, "Only for interface compatible with Gumbel"
        # assert return_logits==False, "Only for interface compatible with Gumbel"
        # reshape z -> (batch, height, width, channel) and flatten
        z = rearrange(z, 'b c h w -> b h w c').contiguous()
        z = l2norm(z)
        z_flattened = z.view(-1, self.embed_dim)
        self.embedding.init_embed_(z_flattened)

        # clculate the distance
        if self.distance == 'l2':
            # l2 distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
            d = - torch.sum(z_flattened.detach() ** 2, dim=1, keepdim=True) - \
                torch.sum(self.embedding.weight ** 2, dim=1) + \
                2 * torch.einsum('bd, dn-> bn', z_flattened.detach(), rearrange(self.embedding.weight, 'n d-> d n'))
        elif self.distance == 'cos':
            # cosine distances from z to embeddings e_j 
            normed_z_flattened = F.normalize(z_flattened, dim=1).detach()
            normed_codebook = F.normalize(self.embedding.weight, dim=1)
            d = torch.einsum('bd,dn->bn', normed_z_flattened, rearrange(normed_codebook, 'n d -> d n'))

        encoding_indices = torch.argmax(d, dim=1)
        z_q = self.embedding(encoding_indices).view(z.shape)
        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)     
        
        if not self.training:
            with torch.no_grad():
                cluster_size = encodings.sum(0)
                self.all_reduce_fn(cluster_size)
                ema_inplace(self.cluster_size, cluster_size, self.decay)


        if self.training and self.embedding.update:

            bins = encodings.sum(0)
            self.all_reduce_fn(bins)

            ema_inplace(self.cluster_size, bins, self.decay)

            zero_mask = (bins == 0)
            bins = bins.masked_fill(zero_mask, 1.)

            embed_sum = z_flattened.t() @ encodings
            self.all_reduce_fn(embed_sum)
                        
            embed_normalized = (embed_sum / bins.unsqueeze(0)).t()
            embed_normalized = l2norm(embed_normalized)
            
            embed_normalized = torch.where(zero_mask[..., None], self.embedding.weight,
                                           embed_normalized)
            norm_ema_inplace(self.embedding.weight, embed_normalized, self.decay)

        # online clustered reinitialisation for unoptimized points
        if self.training:
            avg_probs = torch.mean(encodings, dim=0)
            # calculate the average usage of code entries
            self.embed_prob.mul_(self.decay).add_(avg_probs, alpha= 1 - self.decay)
            # running average updates
            if self.anchor in ['closest', 'random', 'probrandom'] and (not self.init):
                # closest sampling
                if self.anchor == 'closest':
                    sort_distance, indices = d.sort(dim=0)
                    random_feat = z_flattened.detach()[indices[-1,:]]
                # feature pool based random sampling
                elif self.anchor == 'random':
                    random_feat = self.pool.query(z_flattened.detach())
                # probabilitical based random sampling
                elif self.anchor == 'probrandom':
                    norm_distance = F.softmax(d.t(), dim=1)
                    prob = torch.multinomial(norm_distance, num_samples=1).view(-1)
                    random_feat = z_flattened.detach()[prob]
                # decay parameter based on the average usage
                decay = torch.exp(-(self.embed_prob*self.num_embed*10)/(1-self.decay)-1e-3).unsqueeze(1).repeat(1, self.embed_dim)
                self.embedding.weight.data = self.embedding.weight.data * (1 - decay) + random_feat * decay
                if self.first_batch:
                    self.init = True
            # contrastive loss
            if self.contras_loss:
                sort_distance, indices = d.sort(dim=0)
                dis_pos = sort_distance[-max(1, int(sort_distance.size(0)/self.num_embed)):,:].mean(dim=0, keepdim=True)
                dis_neg = sort_distance[:int(sort_distance.size(0)*1/2),:]
                dis = torch.cat([dis_pos, dis_neg], dim=0).t() / 0.07
                contra_loss = F.cross_entropy(dis, torch.zeros((dis.size(0),), dtype=torch.long, device=dis.device))
                loss +=  contra_loss
        # compute loss for embedding
        # loss = self.beta * torch.mean((z_q.detach()-z)**2) + torch.mean((z_q - z.detach()) ** 2)
        loss = self.beta * F.mse_loss(z_q.detach(), z) 

        # preserve gradients
        z_q = z + (z_q - z).detach()
        # reshape back to match original input shape
        z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
        # return z_q, loss, (perplexity, min_encodings, encoding_indices)
        return z_q, loss, encoding_indices

class FeaturePool():
    """
    This class implements a feature buffer that stores previously encoded features
    This buffer enables us to initialize the codebook using a history of generated features
    rather than the ones produced by the latest encoders
    """
    def __init__(self, pool_size, dim=64):
        """
        Initialize the FeaturePool class
        Parameters:
            pool_size(int) -- the size of featue buffer
        """
        self.pool_size = pool_size
        if self.pool_size > 0:
            self.nums_features = 0
            self.features = (torch.rand((pool_size, dim)) * 2 - 1)/ pool_size

    def query(self, features):
        """
        return features from the pool
        """
        self.features = self.features.to(features.device)    
        if self.nums_features < self.pool_size:
            if features.size(0) > self.pool_size: # if the batch size is large enough, directly update the whole codebook
                random_feat_id = torch.randint(0, features.size(0), (int(self.pool_size),))
                self.features = features[random_feat_id]
                self.nums_features = self.pool_size
            else:
                # if the mini-batch is not large nuough, just store it for the next update
                num = self.nums_features + features.size(0)
                self.features[self.nums_features:num] = features
                self.nums_features = num
        else:
            if features.size(0) > int(self.pool_size):
                random_feat_id = torch.randint(0, features.size(0), (int(self.pool_size),))
                self.features = features[random_feat_id]
            else:
                random_id = torch.randperm(self.pool_size)
                self.features[random_id[:features.size(0)]] = features

        return self.features

	# --------------------------------------------------------
	# BEiT v2: Masked Image Modeling with Vector-Quantized Visual Tokenizers (https://arxiv.org/abs/2208.06366)
	# Github source: https://github.com/microsoft/unilm/tree/master/beitv2
	# Copyright (c) 2022 Microsoft
	# Licensed under The MIT License [see LICENSE for details]
	# By Zhiliang Peng
	# Based on VQGAN code bases
	# https://github.com/CompVis/taming-transformers
	# --------------------------------------------------------'

	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	import torch.distributed as distributed
	from einops import rearrange, repeat


	def l2norm(t):
	return F.normalize(t, p = 2, dim = -1)

	def ema_inplace(moving_avg, new, decay):
	moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))

	def sample_vectors(samples, num):
	num_samples, device = samples.shape[0], samples.device

	if num_samples >= num:
	indices = torch.randperm(num_samples, device = device)[:num]
	else:
	indices = torch.randint(0, num_samples, (num,), device = device)

	return samples[indices]

	def kmeans(samples, num_clusters, num_iters = 10, use_cosine_sim = False):

	# rng_state = torch.get_rng_state()
	# torch.manual_seed(42)

	dim, dtype, device = samples.shape[-1], samples.dtype, samples.device

	means = sample_vectors(samples, num_clusters)

	for i in range(num_iters):
	if use_cosine_sim:
	dists = samples @ means.t()
	else:
	diffs = rearrange(samples, 'n d -> n () d') \
	- rearrange(means, 'c d -> () c d')
	dists = -(diffs ** 2).sum(dim = -1)

	buckets = dists.max(dim = -1).indices
	bins = torch.bincount(buckets, minlength = num_clusters)
	zero_mask = bins == 0
	bins_min_clamped = bins.masked_fill(zero_mask, 1)

	new_means = buckets.new_zeros(num_clusters, dim, dtype = dtype)
	new_means.scatter_add_(0, repeat(buckets, 'n -> n d', d = dim), samples)
	new_means = new_means / bins_min_clamped[..., None]

	if use_cosine_sim:
	new_means = l2norm(new_means)

	means = torch.where(zero_mask[..., None], means, new_means)
	print(f"Iteration {i + 1}/{num_iters} - Zero clusters: {zero_mask.sum().item()}")

	# torch.set_rng_state(rng_state)
	return means, bins


	class EmbeddingEMA(nn.Module):
	def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5, kmeans_init=True, codebook_init_path=''):
	super().__init__()
	self.num_tokens = num_tokens
	self.codebook_dim = codebook_dim
	self.decay = decay
	self.eps = eps
	if codebook_init_path == '':
	if not kmeans_init:
	weight = torch.randn(num_tokens, codebook_dim)
	weight = l2norm(weight)
	else:
	weight = torch.zeros(num_tokens, codebook_dim)
	self.register_buffer('initted', torch.Tensor([not kmeans_init]))
	else:
	print(f"load init codebook weight from {codebook_init_path}")
	codebook_ckpt_weight = torch.load(codebook_init_path, map_location='cpu')
	weight = codebook_ckpt_weight.clone()
	self.register_buffer('initted', torch.Tensor([True]))

	self.weight = nn.Parameter(weight, requires_grad = False)
	self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad = False)
	self.embed_avg = nn.Parameter(weight.clone(), requires_grad = False)
	self.update = True

	@torch.jit.ignore
	def init_embed_(self, data):
	if self.initted:
	return
	print("Performing Kemans init for codebook")
	embed, cluster_size = kmeans(data, self.num_tokens, 10, use_cosine_sim = True)
	self.weight.data.copy_(embed)
	self.cluster_size.data.copy_(cluster_size)
	self.initted.data.copy_(torch.Tensor([True]))

	@torch.jit.ignore
	def init_embed_split(self, data, split):
	if self.initted:
	return
	print("Performing Kemans init for codebook")
	embed1, cluster_size1 = kmeans(data[:, :split[0]], self.num_tokens, 10, use_cosine_sim = True)
	embed2, cluster_size2 = kmeans(data[:, split[0]:], self.num_tokens, 10, use_cosine_sim = True)
	embed = torch.cat([embed1, embed2], dim=-1)
	cluster_size = (cluster_size1 + cluster_size2) / 2.
	self.weight.data.copy_(embed)
	self.cluster_size.data.copy_(cluster_size)
	self.initted.data.copy_(torch.Tensor([True]))

	@torch.jit.ignore
	def init_embed_with_ind(self, data, inds):
	if self.initted:
	return
	print("Performing Kemans init for codebook")
	embed, cluster_size = kmeans(data, self.num_tokens, 10, use_cosine_sim = True)
	self.weight.data.copy_(embed)
	self.cluster_size.data.copy_(cluster_size)
	self.initted.data.copy_(torch.Tensor([True]))
	def forward(self, embed_id):
	return F.embedding(embed_id, self.weight)

	def cluster_size_ema_update(self, new_cluster_size):
	self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)

	def embed_avg_ema_update(self, new_embed_avg):
	self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)

	def weight_update(self, num_tokens):
	n = self.cluster_size.sum()
	smoothed_cluster_size = (
	(self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
	)
	#normalize embedding average with smoothed cluster size
	embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
	# embed_normalized = l2norm(self.embed_avg / smoothed_cluster_size.unsqueeze(1))
	self.weight.data.copy_(embed_normalized)

	def norm_ema_inplace(moving_avg, new, decay):
	moving_avg.data.mul_(decay).add_(new, alpha = (1 - decay))
	moving_avg.data.copy_(l2norm(moving_avg.data))

	class NormEMAVectorQuantizer(nn.Module):
	def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
	statistic_code_usage=True, kmeans_init=False, codebook_init_path=''):
	super().__init__()
	self.codebook_dim = embedding_dim
	self.num_tokens = n_embed
	self.beta = beta
	self.decay = decay

	# learnable = True if orthogonal_reg_weight > 0 else False
	self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps, kmeans_init, codebook_init_path)

	self.statistic_code_usage = statistic_code_usage
	if statistic_code_usage:
	self.register_buffer('cluster_size', torch.zeros(n_embed))
	if distributed.is_available() and distributed.is_initialized():
	print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
	self.all_reduce_fn = distributed.all_reduce
	else:
	self.all_reduce_fn = nn.Identity()

	def reset_cluster_size(self, device):
	if self.statistic_code_usage:
	self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
	self.cluster_size = self.cluster_size.to(device)

	def forward(self, z):
	# reshape z -> (batch, height, width, channel) and flatten
	#z, 'b c h w -> b h w c'
	z = rearrange(z, 'b c h w -> b h w c')
	z = l2norm(z)
	z_flattened = z.reshape(-1, self.codebook_dim)

	self.embedding.init_embed_(z_flattened)

	d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \
	self.embedding.weight.pow(2).sum(dim=1) - 2 * \
	torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n'

	encoding_indices = torch.argmin(d, dim=1)

	z_q = self.embedding(encoding_indices).view(z.shape)

	encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)

	if not self.training:
	with torch.no_grad():
	cluster_size = encodings.sum(0)
	self.all_reduce_fn(cluster_size)
	ema_inplace(self.cluster_size, cluster_size, self.decay)

	if self.training and self.embedding.update:
	#EMA cluster size

	bins = encodings.sum(0)
	self.all_reduce_fn(bins)

	ema_inplace(self.cluster_size, bins, self.decay)

	zero_mask = (bins == 0)
	bins = bins.masked_fill(zero_mask, 1.)

	embed_sum = z_flattened.t() @ encodings
	self.all_reduce_fn(embed_sum)

	embed_normalized = (embed_sum / bins.unsqueeze(0)).t()
	embed_normalized = l2norm(embed_normalized)

	embed_normalized = torch.where(zero_mask[..., None], self.embedding.weight,
	embed_normalized)
	norm_ema_inplace(self.embedding.weight, embed_normalized, self.decay)

	loss = self.beta * F.mse_loss(z_q.detach(), z)

	z_q = z + (z_q - z).detach()

	#z_q, 'b h w c -> b c h w'
	z_q = rearrange(z_q, 'b h w c -> b c h w')
	return z_q, loss, encoding_indices



	class VectorQuantizer(nn.Module):
	def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
	statistic_code_usage=True, kmeans_init=False, codebook_init_path=''):
	super().__init__()
	self.codebook_dim = embedding_dim
	self.num_tokens = n_embed
	self.beta = beta
	decay = 0.
	self.decay = decay

	# learnable = True if orthogonal_reg_weight > 0 else False
	self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps, kmeans_init, codebook_init_path)

	self.statistic_code_usage = statistic_code_usage
	if statistic_code_usage:
	self.register_buffer('cluster_size', torch.zeros(n_embed))
	if distributed.is_available() and distributed.is_initialized():
	print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
	self.all_reduce_fn = distributed.all_reduce
	else:
	self.all_reduce_fn = nn.Identity()

	def reset_cluster_size(self, device):
	if self.statistic_code_usage:
	self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
	self.cluster_size = self.cluster_size.to(device)

	def forward(self, bhw, encoding_indices, z=None):
	embed_size = [bhw[0], bhw[1], bhw[2], -1]
	z_q = self.embedding(encoding_indices).view(embed_size)

	encodings = F.one_hot(encoding_indices, self.num_tokens).type(z_q.dtype)

	if not self.training:
	with torch.no_grad():
	cluster_size = encodings.sum(0)
	self.all_reduce_fn(cluster_size)
	ema_inplace(self.cluster_size, cluster_size, self.decay)
	if z is not None:
	#EMA cluster size
	z = rearrange(z, 'b c h w -> b h w c')
	loss = self.beta * F.mse_loss(z_q, z.detach())

	else:
	loss = torch.tensor([0.]).mean().to(z_q)
	z_q = rearrange(z_q, 'b h w c -> b c h w')
	return z_q, loss, encoding_indices



	class CVectorQuantiser(nn.Module):
	"""
	Improved version over vector quantiser, with the dynamic initialisation
	for these unoptimised "dead" points.
	num_embed: number of codebook entry
	embed_dim: dimensionality of codebook entry
	beta: weight for the commitment loss
	distance: distance for looking up the closest code
	anchor: anchor sampled methods
	first_batch: if true, the offline version of our model
	contras_loss: if true, use the contras_loss to further improve the performance
	"""
	def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
	statistic_code_usage=True, kmeans_init=True, codebook_init_path='',
	distance='l2', anchor='closest', first_batch=False, contras_loss=False):
	super().__init__()

	self.num_embed = n_embed
	self.n_e = n_embed
	self.num_tokens = n_embed
	self.embed_dim = embedding_dim

	self.beta = beta
	self.distance = distance
	self.anchor = anchor
	self.first_batch = first_batch
	self.contras_loss = contras_loss
	self.decay = decay
	self.init = False

	self.pool = FeaturePool(self.num_embed, self.embed_dim)
	# self.embedding = nn.Embedding(self.num_embed, self.embed_dim)
	# self.embedding.weight.data.uniform_(-1.0 / self.num_embed, 1.0 / self.num_embed)
	self.embedding = EmbeddingEMA(self.num_tokens, self.embed_dim, decay, eps, kmeans_init, codebook_init_path)

	self.register_buffer("embed_prob", torch.zeros(self.num_embed))
	self.statistic_code_usage = statistic_code_usage
	if statistic_code_usage:
	self.register_buffer('cluster_size', torch.zeros(n_embed))
	if distributed.is_available() and distributed.is_initialized():
	print("ddp is enable, so use ddp_reduce to sync the statistic_code_usage for each gpu!")
	self.all_reduce_fn = distributed.all_reduce
	else:
	self.all_reduce_fn = nn.Identity()

	def reset_cluster_size(self, device):
	if self.statistic_code_usage:
	self.register_buffer('cluster_size', torch.zeros(self.num_tokens))
	self.cluster_size = self.cluster_size.to(device)

	def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
	# assert temp is None or temp==1.0, "Only for interface compatible with Gumbel"
	# assert rescale_logits==False, "Only for interface compatible with Gumbel"
	# assert return_logits==False, "Only for interface compatible with Gumbel"
	# reshape z -> (batch, height, width, channel) and flatten
	z = rearrange(z, 'b c h w -> b h w c').contiguous()
	z = l2norm(z)
	z_flattened = z.view(-1, self.embed_dim)
	self.embedding.init_embed_(z_flattened)

	# clculate the distance
	if self.distance == 'l2':
	# l2 distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
	d = - torch.sum(z_flattened.detach() ** 2, dim=1, keepdim=True) - \
	torch.sum(self.embedding.weight ** 2, dim=1) + \
	2 * torch.einsum('bd, dn-> bn', z_flattened.detach(), rearrange(self.embedding.weight, 'n d-> d n'))
	elif self.distance == 'cos':
	# cosine distances from z to embeddings e_j
	normed_z_flattened = F.normalize(z_flattened, dim=1).detach()
	normed_codebook = F.normalize(self.embedding.weight, dim=1)
	d = torch.einsum('bd,dn->bn', normed_z_flattened, rearrange(normed_codebook, 'n d -> d n'))

	encoding_indices = torch.argmax(d, dim=1)
	z_q = self.embedding(encoding_indices).view(z.shape)
	encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)

	if not self.training:
	with torch.no_grad():
	cluster_size = encodings.sum(0)
	self.all_reduce_fn(cluster_size)
	ema_inplace(self.cluster_size, cluster_size, self.decay)


	if self.training and self.embedding.update:

	bins = encodings.sum(0)
	self.all_reduce_fn(bins)

	ema_inplace(self.cluster_size, bins, self.decay)

	zero_mask = (bins == 0)
	bins = bins.masked_fill(zero_mask, 1.)

	embed_sum = z_flattened.t() @ encodings
	self.all_reduce_fn(embed_sum)

	embed_normalized = (embed_sum / bins.unsqueeze(0)).t()
	embed_normalized = l2norm(embed_normalized)

	embed_normalized = torch.where(zero_mask[..., None], self.embedding.weight,
	embed_normalized)
	norm_ema_inplace(self.embedding.weight, embed_normalized, self.decay)

	# online clustered reinitialisation for unoptimized points
	if self.training:
	avg_probs = torch.mean(encodings, dim=0)
	# calculate the average usage of code entries
	self.embed_prob.mul_(self.decay).add_(avg_probs, alpha= 1 - self.decay)
	# running average updates
	if self.anchor in ['closest', 'random', 'probrandom'] and (not self.init):
	# closest sampling
	if self.anchor == 'closest':
	sort_distance, indices = d.sort(dim=0)
	random_feat = z_flattened.detach()[indices[-1,:]]
	# feature pool based random sampling
	elif self.anchor == 'random':
	random_feat = self.pool.query(z_flattened.detach())
	# probabilitical based random sampling
	elif self.anchor == 'probrandom':
	norm_distance = F.softmax(d.t(), dim=1)
	prob = torch.multinomial(norm_distance, num_samples=1).view(-1)
	random_feat = z_flattened.detach()[prob]
	# decay parameter based on the average usage
	decay = torch.exp(-(self.embed_probself.num_embed10)/(1-self.decay)-1e-3).unsqueeze(1).repeat(1, self.embed_dim)
	self.embedding.weight.data = self.embedding.weight.data * (1 - decay) + random_feat * decay
	if self.first_batch:
	self.init = True
	# contrastive loss
	if self.contras_loss:
	sort_distance, indices = d.sort(dim=0)
	dis_pos = sort_distance[-max(1, int(sort_distance.size(0)/self.num_embed)):,:].mean(dim=0, keepdim=True)
	dis_neg = sort_distance[:int(sort_distance.size(0)*1/2),:]
	dis = torch.cat([dis_pos, dis_neg], dim=0).t() / 0.07
	contra_loss = F.cross_entropy(dis, torch.zeros((dis.size(0),), dtype=torch.long, device=dis.device))
	loss += contra_loss
	# compute loss for embedding
	# loss = self.beta * torch.mean((z_q.detach()-z)2) + torch.mean((z_q - z.detach()) 2)
	loss = self.beta * F.mse_loss(z_q.detach(), z)

	# preserve gradients
	z_q = z + (z_q - z).detach()
	# reshape back to match original input shape
	z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
	# return z_q, loss, (perplexity, min_encodings, encoding_indices)
	return z_q, loss, encoding_indices

	class FeaturePool():
	"""
	This class implements a feature buffer that stores previously encoded features

	This buffer enables us to initialize the codebook using a history of generated features
	rather than the ones produced by the latest encoders
	"""
	def __init__(self, pool_size, dim=64):
	"""
	Initialize the FeaturePool class

	Parameters:
	pool_size(int) -- the size of featue buffer
	"""
	self.pool_size = pool_size
	if self.pool_size > 0:
	self.nums_features = 0
	self.features = (torch.rand((pool_size, dim)) * 2 - 1)/ pool_size

	def query(self, features):
	"""
	return features from the pool
	"""
	self.features = self.features.to(features.device)
	if self.nums_features < self.pool_size:
	if features.size(0) > self.pool_size: # if the batch size is large enough, directly update the whole codebook
	random_feat_id = torch.randint(0, features.size(0), (int(self.pool_size),))
	self.features = features[random_feat_id]
	self.nums_features = self.pool_size
	else:
	# if the mini-batch is not large nuough, just store it for the next update
	num = self.nums_features + features.size(0)
	self.features[self.nums_features:num] = features
	self.nums_features = num
	else:
	if features.size(0) > int(self.pool_size):
	random_feat_id = torch.randint(0, features.size(0), (int(self.pool_size),))
	self.features = features[random_feat_id]
	else:
	random_id = torch.randperm(self.pool_size)
	self.features[random_id[:features.size(0)]] = features

	return self.features