|
|
|
|
|
|
|
|
|
|
|
import math |
|
import torch |
|
from torch import nn |
|
from .fvq import FactorizedVectorQuantize |
|
|
|
|
|
class ResidualVQ(nn.Module): |
|
"""Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf""" |
|
|
|
def __init__(self, *, num_quantizers, codebook_size, **kwargs): |
|
super().__init__() |
|
VQ = FactorizedVectorQuantize |
|
if type(codebook_size) == int: |
|
codebook_size = [codebook_size] * num_quantizers |
|
self.layers = nn.ModuleList( |
|
[VQ(codebook_size=2**size, **kwargs) for size in codebook_size] |
|
) |
|
self.num_quantizers = num_quantizers |
|
self.quantizer_dropout = kwargs.get("quantizer_dropout", 0.0) |
|
self.dropout_type = kwargs.get("dropout_type", None) |
|
|
|
def forward(self, x, n_quantizers=None): |
|
quantized_out = 0.0 |
|
residual = x |
|
|
|
all_losses = [] |
|
all_indices = [] |
|
all_quantized = [] |
|
|
|
if n_quantizers is None: |
|
n_quantizers = self.num_quantizers |
|
if self.training: |
|
n_quantizers = torch.ones((x.shape[0],)) * self.num_quantizers + 1 |
|
if self.dropout_type == "linear": |
|
dropout = torch.randint(1, self.num_quantizers + 1, (x.shape[0],)) |
|
elif self.dropout_type == "exp": |
|
dropout = torch.randint( |
|
1, int(math.log2(self.num_quantizers)), (x.shape[0],) |
|
) |
|
dropout = torch.pow(2, dropout) |
|
n_dropout = int(x.shape[0] * self.quantizer_dropout) |
|
n_quantizers[:n_dropout] = dropout[:n_dropout] |
|
n_quantizers = n_quantizers.to(x.device) |
|
|
|
for idx, layer in enumerate(self.layers): |
|
if not self.training and idx >= n_quantizers: |
|
break |
|
quantized, indices, loss = layer(residual) |
|
|
|
mask = ( |
|
torch.full((x.shape[0],), fill_value=idx, device=x.device) |
|
< n_quantizers |
|
) |
|
|
|
residual = residual - quantized |
|
|
|
quantized_out = quantized_out + quantized * mask[:, None, None] |
|
|
|
|
|
loss = (loss * mask).mean() |
|
|
|
all_indices.append(indices) |
|
all_losses.append(loss) |
|
all_quantized.append(quantized) |
|
all_losses, all_indices, all_quantized = map( |
|
torch.stack, (all_losses, all_indices, all_quantized) |
|
) |
|
return quantized_out, all_indices, all_losses, all_quantized |
|
|
|
def vq2emb(self, vq): |
|
|
|
quantized_out = 0.0 |
|
for idx, layer in enumerate(self.layers): |
|
quantized = layer.vq2emb(vq[idx]) |
|
quantized_out += quantized |
|
return quantized_out |
|
|
|
def get_emb(self): |
|
embs = [] |
|
for idx, layer in enumerate(self.layers): |
|
embs.append(layer.get_emb()) |
|
return embs |
|
|
