402 lines
13 KiB
Python
402 lines
13 KiB
Python
# Copyright (c) 2024 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from einops import rearrange, repeat
|
|
from torch.nn.utils import weight_norm
|
|
|
|
|
|
def WNConv1d(*args, **kwargs):
|
|
return weight_norm(nn.Conv1d(*args, **kwargs))
|
|
|
|
|
|
def WNConvTranspose1d(*args, **kwargs):
|
|
return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
|
|
|
|
|
|
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 laplace_smoothing(x, n_categories, eps=1e-5):
|
|
return (x + eps) / (x.sum() + n_categories * eps)
|
|
|
|
|
|
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):
|
|
dim, dtype, device = samples.shape[-1], samples.dtype, samples.device
|
|
|
|
means = sample_vectors(samples, num_clusters)
|
|
|
|
for _ 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)
|
|
|
|
return means, bins
|
|
|
|
|
|
class EuclideanCodebook(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
codebook_size,
|
|
kmeans_init=False,
|
|
kmeans_iters=10,
|
|
decay=0.8,
|
|
eps=1e-5,
|
|
threshold_ema_dead_code=2,
|
|
weight_init=False,
|
|
):
|
|
super().__init__()
|
|
|
|
self.decay = decay
|
|
init_fn = torch.randn if not weight_init else torch.zeros
|
|
embed = init_fn(codebook_size, dim)
|
|
|
|
if weight_init:
|
|
nn.init.uniform_(embed, -1 / codebook_size, 1 / codebook_size)
|
|
|
|
self.codebook_size = codebook_size
|
|
self.kmeans_iters = kmeans_iters
|
|
self.eps = eps
|
|
self.threshold_ema_dead_code = threshold_ema_dead_code
|
|
|
|
self.register_buffer(
|
|
"initted", torch.Tensor([not kmeans_init])
|
|
) # if kmeans_init is True, then initted is False; otherwise, initted is True
|
|
self.register_buffer("cluster_size", torch.zeros(codebook_size))
|
|
self.register_buffer("embed", embed)
|
|
self.register_buffer("embed_avg", embed.clone())
|
|
|
|
def init_embed_(self, data):
|
|
embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
|
|
self.embed.data.copy_(embed)
|
|
self.embed_avg.data.copy_(embed)
|
|
self.cluster_size.data.copy_(cluster_size)
|
|
self.initted.data.copy_(torch.Tensor([True]))
|
|
|
|
def replace(self, samples, mask):
|
|
modified_codebook = torch.where(
|
|
mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
|
|
)
|
|
self.embed.data.copy_(modified_codebook)
|
|
|
|
def expire_codes_(self, batch_samples):
|
|
if self.threshold_ema_dead_code == 0:
|
|
return
|
|
|
|
expired_codes = self.cluster_size < self.threshold_ema_dead_code
|
|
if not torch.any(expired_codes):
|
|
return
|
|
batch_samples = rearrange(batch_samples, "... d -> (...) d")
|
|
self.replace(batch_samples, mask=expired_codes)
|
|
|
|
def forward(self, x):
|
|
shape, dtype = x.shape, x.dtype
|
|
flatten = rearrange(x, "... d -> (...) d")
|
|
embed = self.embed.t() # (codebook_size, dim) -> (dim, codebook_size)
|
|
|
|
if not self.initted:
|
|
self.init_embed_(flatten)
|
|
|
|
dist = -(
|
|
flatten.pow(2).sum(1, keepdim=True)
|
|
- 2 * flatten @ embed
|
|
+ embed.pow(2).sum(0, keepdim=True)
|
|
)
|
|
|
|
embed_ind = dist.max(dim=-1).indices
|
|
embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
|
|
embed_ind = embed_ind.view(*shape[:-1])
|
|
quantize = F.embedding(embed_ind, self.embed)
|
|
|
|
if self.training:
|
|
ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
|
|
embed_sum = (
|
|
flatten.t() @ embed_onehot
|
|
) # (dim, ...) @ (..., codebook_size) -> (dim, codebook_size)
|
|
ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
|
|
cluster_size = (
|
|
laplace_smoothing(self.cluster_size, self.codebook_size, self.eps)
|
|
* self.cluster_size.sum()
|
|
)
|
|
embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
|
|
self.embed.data.copy_(embed_normalized)
|
|
self.expire_codes_(x)
|
|
|
|
return quantize, embed_ind
|
|
|
|
def vq2emb(self, vq):
|
|
quantize = F.embedding(vq, self.embed)
|
|
return quantize
|
|
|
|
def latent2dist(self, x):
|
|
shape, dtype = x.shape, x.dtype
|
|
flatten = rearrange(x, "... d -> (...) d")
|
|
embed = self.embed.t() # (codebook_size, dim) -> (dim, codebook_size)
|
|
|
|
if not self.initted:
|
|
self.init_embed_(flatten)
|
|
|
|
dist = -(
|
|
flatten.pow(2).sum(1, keepdim=True)
|
|
- 2 * flatten @ embed
|
|
+ embed.pow(2).sum(0, keepdim=True)
|
|
)
|
|
|
|
embed_ind = dist.max(dim=-1).indices
|
|
embed_ind = embed_ind.view(*shape[:-1])
|
|
quantize = F.embedding(embed_ind, self.embed)
|
|
|
|
dist = dist.view(*shape[:-1], -1)
|
|
|
|
return dist, embed_ind, quantize
|
|
|
|
|
|
class SimpleCodebook(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
codebook_size,
|
|
use_l2_normlize=False,
|
|
):
|
|
super().__init__()
|
|
|
|
self.dim = dim
|
|
self.codebook_size = codebook_size
|
|
self.use_l2_normlize = use_l2_normlize
|
|
|
|
self.embed = nn.Embedding(self.codebook_size, self.dim)
|
|
|
|
def forward(self, x):
|
|
shape, dtype = x.shape, x.dtype
|
|
flatten = rearrange(x, "... d -> (...) d")
|
|
embed = self.embed.weight.t() # (codebook_size, dim) -> (dim, codebook_size)
|
|
|
|
if self.use_l2_normlize:
|
|
flatten = F.normalize(flatten)
|
|
embed = F.normalize(embed)
|
|
|
|
dist = -(
|
|
flatten.pow(2).sum(1, keepdim=True)
|
|
- 2 * flatten @ embed
|
|
+ embed.pow(2).sum(0, keepdim=True)
|
|
)
|
|
|
|
embed_ind = dist.max(dim=-1).indices
|
|
embed_ind = embed_ind.view(*shape[:-1])
|
|
quantize = F.embedding(embed_ind, self.embed)
|
|
|
|
return quantize, embed_ind
|
|
|
|
def vq2emb(self, vq):
|
|
quantize = F.embedding(vq, self.embed.weight)
|
|
return quantize
|
|
|
|
def latent2dist(self, x):
|
|
shape, dtype = x.shape, x.dtype
|
|
flatten = rearrange(x, "... d -> (...) d")
|
|
embed = self.embed.weight.t() # (codebook_size, dim) -> (dim, codebook_size)
|
|
|
|
if self.use_l2_normlize:
|
|
flatten = F.normalize(flatten)
|
|
embed = F.normalize(embed)
|
|
|
|
dist = -(
|
|
flatten.pow(2).sum(1, keepdim=True)
|
|
- 2 * flatten @ embed
|
|
+ embed.pow(2).sum(0, keepdim=True)
|
|
)
|
|
|
|
embed_ind = dist.max(dim=-1).indices
|
|
embed_ind = embed_ind.view(*shape[:-1])
|
|
quantize = F.embedding(embed_ind, self.embed)
|
|
|
|
dist = dist.view(*shape[:-1], -1)
|
|
|
|
return dist, embed_ind, quantize
|
|
|
|
|
|
class VectorQuantize(nn.Module):
|
|
"""Vector quantization and factorized vecotor quantization implementation
|
|
Args:
|
|
input_dim (int): Dimension of input.
|
|
codebook_size (int): Codebook size.
|
|
codebook_dim (int): Codebook dimension. We suggest use codebook_dim = input_dim
|
|
if use codebook_type == "euclidean", otherwise, if you want to use
|
|
factorized vector quantization, use codebook_dim as small number (e.g. 8 or 32).
|
|
commitment (float): Weight for commitment loss.
|
|
use_l2_normlize (bool): Whether to use l2 normlized codes for factorized vecotor quantization,
|
|
we suggest use it as True if you want to use factorized vector quantization
|
|
kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
|
|
kmeans_iters (int): Number of iterations used for kmeans initialization.
|
|
decay (float): Decay for exponential moving average over the codebooks.
|
|
epsilon (float): Epsilon value for numerical stability.
|
|
threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
|
|
that have an exponential moving average cluster size less than the specified threshold with
|
|
randomly selected vector from the current batch.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_dim,
|
|
codebook_size,
|
|
codebook_dim,
|
|
commitment=0.005,
|
|
codebook_loss_weight=1.0,
|
|
use_l2_normlize=False,
|
|
codebook_type="euclidean", # "euclidean" or "simple"
|
|
kmeans_init=False,
|
|
kmeans_iters=10,
|
|
decay=0.8,
|
|
eps=1e-5,
|
|
threshold_ema_dead_code=2,
|
|
weight_init=False,
|
|
):
|
|
super().__init__()
|
|
self.input_dim = input_dim
|
|
self.codebook_size = codebook_size
|
|
self.codebook_dim = codebook_dim
|
|
self.commitment = commitment
|
|
self.codebook_loss_weight = codebook_loss_weight
|
|
self.use_l2_normlize = use_l2_normlize
|
|
self.codebook_type = codebook_type
|
|
self.kmeans_init = kmeans_init
|
|
self.kmeans_iters = kmeans_iters
|
|
self.decay = decay
|
|
self.eps = eps
|
|
self.threshold_ema_dead_code = threshold_ema_dead_code
|
|
self.weight_init = weight_init
|
|
|
|
if self.input_dim != self.codebook_dim:
|
|
self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
|
|
self.out_project = WNConv1d(
|
|
self.codebook_dim, self.input_dim, kernel_size=1
|
|
)
|
|
|
|
else:
|
|
self.in_project = nn.Identity()
|
|
self.out_project = nn.Identity()
|
|
|
|
if self.codebook_type == "euclidean":
|
|
self.codebook = EuclideanCodebook(
|
|
self.codebook_dim,
|
|
codebook_size=self.codebook_size,
|
|
kmeans_init=self.kmeans_init,
|
|
kmeans_iters=self.kmeans_iters,
|
|
decay=self.decay,
|
|
eps=self.eps,
|
|
threshold_ema_dead_code=self.threshold_ema_dead_code,
|
|
weight_init=self.weight_init,
|
|
)
|
|
elif self.codebook_type == "simple":
|
|
self.codebook = SimpleCodebook(
|
|
self.codebook_dim,
|
|
codebook_size=self.codebook_size,
|
|
use_l2_normlize=self.use_l2_normlize,
|
|
)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"codebook_type {self.codebook_type} is not implemented!"
|
|
)
|
|
|
|
def forward(self, z):
|
|
"""
|
|
Parameters
|
|
----------
|
|
z: torch.Tensor[B x D x T]
|
|
|
|
Returns
|
|
-------
|
|
z_q: torch.Tensor[B x D x T]
|
|
Quantized continuous representation of input
|
|
commit_loss: Tensor[B]
|
|
Commitment loss to train encoder to predict vectors closer to codebook entries
|
|
codebook_loss: Tensor[B]
|
|
Codebook loss to update the codebook
|
|
indices: torch.Tensor[B x T]
|
|
Codebook indices (quantized discrete representation of input)
|
|
z_e: torch.Tensor[B x D x T]
|
|
Projected latents (continuous representation of input before quantization)
|
|
"""
|
|
|
|
# Factorized codes project input into low-dimensional space if self.input_dim != self.codebook_dim
|
|
z_e = self.in_project(z)
|
|
z_q, indices = self.decode_latents(z_e)
|
|
|
|
# Compute commitment loss and codebook loss
|
|
if self.training:
|
|
commit_loss = (
|
|
F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
|
|
* self.commitment
|
|
)
|
|
codebook_loss = (
|
|
F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
|
|
* self.codebook_loss_weight
|
|
)
|
|
else:
|
|
commit_loss = torch.zeros(z.shape[0], device=z.device)
|
|
codebook_loss = torch.zeros(z.shape[0], device=z.device)
|
|
|
|
z_q = z_e + (z_q - z_e).detach()
|
|
|
|
z_q = self.out_project(z_q)
|
|
|
|
return z_q, commit_loss, codebook_loss, indices, z_e
|
|
|
|
def decode_latents(self, latents):
|
|
encodings = rearrange(latents, "b d t -> b t d")
|
|
z_q, indices = self.codebook(encodings)
|
|
z_q = z_q.transpose(1, 2)
|
|
return z_q, indices
|
|
|
|
def vq2emb(self, vq, out_proj=True):
|
|
emb = self.codebook.vq2emb(vq)
|
|
emb = emb.transpose(1, 2)
|
|
if out_proj:
|
|
emb = self.out_project(emb)
|
|
return emb
|
|
|
|
def latent2dist(self, latents):
|
|
latents = rearrange(latents, "b d t -> b t d")
|
|
dist, embed_ind, quantize = self.codebook.latent2dist(latents)
|
|
return dist, embed_ind, quantize.transpose(1, 2)
|