add EMA Vector Quantizer

Author: tgisaturday

taming/models/vqgan.py

@@ -7,7 +7,7 @@
 from taming.modules.diffusionmodules.model import Encoder, Decoder
 from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
 from taming.modules.vqvae.quantize import GumbelQuantize
-
+from taming.modules.vqvae.quantize import EMAVectorQuantizer
 
 class VQModel(pl.LightningModule):
     def __init__(self,
@@ -361,3 +361,44 @@ def log_images(self, batch, **kwargs):
         log["inputs"] = x
         log["reconstructions"] = x_rec
         return log
+
+
+class EMAVQ(VQModel):
+    def __init__(self,
+                 ddconfig,
+                 lossconfig,
+                 n_embed,
+                 embed_dim,
+                 ckpt_path=None,
+                 ignore_keys=[],
+                 image_key="image",
+                 colorize_nlabels=None,
+                 monitor=None,
+                 remap=None,
+                 sane_index_shape=False,  # tell vector quantizer to return indices as bhw
+                 ):
+        super().__init__(ddconfig,
+                         lossconfig,
+                         n_embed,
+                         embed_dim,
+                         ckpt_path=None,
+                         ignore_keys=ignore_keys,
+                         image_key=image_key,
+                         colorize_nlabels=colorize_nlabels,
+                         monitor=monitor,
+                         )
+        self.quantize = EMAVectorQuantizer(n_embed=n_embed,
+                                           embedding_dim=embed_dim,
+                                           beta=0.25,
+                                           remap=remap)
+    def configure_optimizers(self):
+        lr = self.learning_rate
+        #Remove self.quantize from parameter list since it is updated via EMA
+        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
+                                  list(self.decoder.parameters())+
+                                  list(self.quant_conv.parameters())+
+                                  list(self.post_quant_conv.parameters()),
+                                  lr=lr, betas=(0.5, 0.9))
+        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+                                    lr=lr, betas=(0.5, 0.9))
+        return [opt_ae, opt_disc], []                                           

taming/modules/vqvae/quantize.py

@@ -327,3 +327,165 @@ def get_codebook_entry(self, indices, shape):
             z_q = z_q.permute(0, 3, 1, 2).contiguous()
 
         return z_q
+
+
+
+class EMAVectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
+                remap=None, unknown_index="random"):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.n_embed = n_embed
+        self.decay = decay
+        self.eps = eps
+        self.beta = beta
+        self.embedding = nn.Embedding(self.n_embed, self.embedding_dim)
+        self.embedding.weight.requires_grad = False
+        self.cluster_size = nn.Parameter(torch.zeros(n_embed),requires_grad=False)
+        self.embed_avg = nn.Parameter(torch.Tensor(self.n_embed, self.embedding_dim),requires_grad=False)
+        self.embed_avg.data.copy_(self.embedding.weight.data)
+        self.remap = remap
+        if self.remap is not None:
+            self.register_buffer("used", torch.tensor(np.load(self.remap)))
+            self.re_embed = self.used.shape[0]
+            self.unknown_index = unknown_index # "random" or "extra" or integer
+            if self.unknown_index == "extra":
+                self.unknown_index = self.re_embed
+                self.re_embed = self.re_embed+1
+            print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
+                  f"Using {self.unknown_index} for unknown indices.")
+        else:
+            self.re_embed = n_embed
+
+    def remap_to_used(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        match = (inds[:,:,None]==used[None,None,...]).long()
+        new = match.argmax(-1)
+        unknown = match.sum(2)<1
+        if self.unknown_index == "random":
+            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
+        else:
+            new[unknown] = self.unknown_index
+        return new.reshape(ishape)
+
+    def unmap_to_all(self, inds):
+        ishape = inds.shape
+        assert len(ishape)>1
+        inds = inds.reshape(ishape[0],-1)
+        used = self.used.to(inds)
+        if self.re_embed > self.used.shape[0]: # extra token
+            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
+        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
+        return back.reshape(ishape)
+
+    def forward(self, z):
+        # reshape z -> (batch, height, width, channel) and flatten
+        #z, 'b c h w -> b h w c'
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.view(-1, self.embedding_dim)
+        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+
+        d = torch.sum(z_flattened.pow(2), dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight.pow(2), dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', z_flattened, self.embedding.weight.permute(1,0)) # '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.n_embed).type(z.dtype)     
+        avg_probs = torch.mean(encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        if self.training:
+            encodings_sum = encodings.sum(0)
+            #EMA cluster size
+            self.cluster_size.mul_(self.decay).add_(encodings_sum, alpha=1 - self.decay)
+
+            embed_sum = torch.matmul(encodings.t(), z_flattened)
+            #EMA embedding average
+            self.embed_avg.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay)
+
+            #cluster size Laplace smoothing 
+            n = self.cluster_size.sum()
+            cluster_size = (
+                (self.cluster_size + self.eps) / (n + self.n_embed * self.eps) * n
+            )
+            #normalize embedding average with smoothed cluster size
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embedding.weight.data.copy_(embed_normalized.data)
+
+        # compute loss for embedding
+        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, 'b h w c -> b c h w'
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+        return z_q, loss, (perplexity, encodings, encoding_indices)
+
+
+
+#Original Sonnet version of EMAVectorQuantizer
+class EmbeddingEMA(nn.Module):
+    def __init__(self, n_embed, embedding_dim):
+        super().__init__()
+        weight = torch.randn(embedding_dim, n_embed)
+        self.register_buffer("weight", weight)
+        self.register_buffer("cluster_size", torch.zeros(n_embed))
+        self.register_buffer("embed_avg", weight.clone())
+
+    def forward(self, embed_id):
+        return F.embedding(embed_id, self.weight.transpose(0, 1))
+
+
+class SonnetEMAVectorQuantizer(nn.Module):
+    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
+                remap=None, unknown_index="random"):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.n_embed = n_embed
+        self.decay = decay
+        self.eps = eps
+        self.beta = beta
+        self.embedding = EmbeddingEMA(n_embed,embedding_dim)
+
+    def forward(self, z):
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.reshape(-1, self.embedding_dim)
+        d = (
+            z_flattened.pow(2).sum(1, keepdim=True)
+            - 2 * z_flattened @ self.embedding.weight
+            + self.embedding.weight.pow(2).sum(0, keepdim=True)
+        )
+        _, encoding_indices = (-d).max(1)
+        encodings = F.one_hot(encoding_indices, self.n_embed).type(z_flattened.dtype)
+        encoding_indices = encoding_indices.view(*z.shape[:-1])
+        z_q = self.embedding(encoding_indices)
+        avg_probs = torch.mean(encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        if self.training:
+            encodings_sum = encodings.sum(0)
+            embed_sum = z_flattened.transpose(0, 1) @ encodings
+            #EMA cluster size
+            self.embedding.cluster_size.data.mul_(self.decay).add_(encodings_sum, alpha=1 - self.decay)
+            #EMA embedding average
+            self.embedding.embed_avg.data.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay)
+
+            #cluster size Laplace smoothing 
+            n = self.embedding.cluster_size.sum()
+            cluster_size = (
+                (self.embedding.cluster_size + self.eps) / (n + self.n_embed * self.eps) * n
+            )
+            #normalize embedding average with smoothed cluster size
+            embed_normalized = self.embedding.embed_avg / cluster_size.unsqueeze(0)
+            self.embedding.weight.data.copy_(embed_normalized)
+
+        loss = self.beta * (z_q.detach() - z).pow(2).mean()
+        z_q = z + (z_q - z).detach()
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+        return z_q, loss, (perplexity, encodings, encoding_indices)