# -------------------------------------------------------- # 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