| | |
| | |
| | |
| | |
| |
|
| | import torch |
| | import torch.nn as nn |
| |
|
| | from fairseq.modules import Fp32GroupNorm |
| |
|
| |
|
| | class KmeansVectorQuantizer(nn.Module): |
| | def __init__( |
| | self, dim, num_vars, groups, combine_groups, vq_dim, time_first, gamma=0.25 |
| | ): |
| | '''Vector quantization using straight pass-through estimator (i.e. kmeans) |
| | |
| | Args: |
| | dim: input dimension (channels) |
| | num_vars: number of quantized vectors per group |
| | groups: number of groups for vector quantization |
| | combine_groups: whether to use the vectors for all groups |
| | vq_dim: dimensionality of the resulting quantized vector |
| | time_first: if true, expect input in BxTxC format, otherwise in BxCxT |
| | gamma: commitment loss coefficient |
| | ''' |
| | super().__init__() |
| |
|
| | self.groups = groups |
| | self.combine_groups = combine_groups |
| | self.input_dim = dim |
| | self.num_vars = num_vars |
| | self.vq_dim = vq_dim |
| | self.time_first = time_first |
| |
|
| | assert ( |
| | vq_dim % groups == 0 |
| | ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation" |
| |
|
| | self.var_dim = vq_dim // groups |
| | num_groups = groups if not combine_groups else 1 |
| |
|
| | self.embedding = nn.Parameter( |
| | 0.01 * torch.randn(num_vars, num_groups, self.var_dim) |
| | ) |
| | self.projection = nn.Sequential( |
| | nn.Conv1d(dim, dim, kernel_size=1, groups=groups, bias=False), |
| | Fp32GroupNorm(groups, dim), |
| | ) |
| | self.gamma = gamma |
| | self.mse_mean = nn.MSELoss(reduction="mean") |
| |
|
| | def _pass_grad(self, x, y): |
| | """ Manually set gradient for backward pass. |
| | for y = f(x), ensure that during the backward pass, |
| | dL/dy = dL/dx regardless of f(x). |
| | Returns: |
| | y, with the gradient forced to be dL/dy = dL/dx. |
| | """ |
| |
|
| | return y.detach() + (x - x.detach()) |
| |
|
| | @property |
| | def expand_embedding(self): |
| | if self.combine_groups: |
| | return self.embedding.expand(self.num_vars, self.groups, self.var_dim) |
| | return self.embedding |
| |
|
| | def forward_idx(self, x): |
| | res = self.forward(x, produce_targets=True) |
| | return res["x"], res["targets"] |
| |
|
| | def forward(self, x, produce_targets=False): |
| |
|
| | result = {"num_vars": self.num_vars} |
| |
|
| | if self.time_first: |
| | x = x.transpose(1, 2) |
| |
|
| | bsz, fsz, tsz = x.shape |
| |
|
| | ze = self.projection(x) |
| | ze_ = ze.view(bsz, self.groups, self.var_dim, tsz).permute(0, 3, 1, 2) |
| | d = ( |
| | (ze_.unsqueeze(0) - self.expand_embedding.unsqueeze(1).unsqueeze(1)) |
| | .view(self.num_vars, bsz, tsz, self.groups, -1) |
| | .norm(dim=-1, p=2) |
| | ) |
| | idx = d.argmin(dim=0) |
| | zq = ( |
| | torch.stack( |
| | [ |
| | self.expand_embedding[idx[..., group], group] |
| | for group in range(self.groups) |
| | ], |
| | dim=-2, |
| | ) |
| | .view(bsz, tsz, self.groups * self.var_dim) |
| | .permute(0, 2, 1) |
| | ) |
| | assert ze.shape == zq.shape, (ze.shape, zq.shape) |
| | x = self._pass_grad(ze, zq) |
| |
|
| | hard_x = ( |
| | idx.new_zeros(bsz*tsz*self.groups, self.num_vars) |
| | .scatter_(-1, idx.view(-1, 1), 1.0) |
| | .view(bsz * tsz, self.groups, -1) |
| | ) |
| | hard_probs = torch.mean(hard_x.float(), dim=0) |
| | result["code_perplexity"] = torch.exp( |
| | -torch.sum(hard_probs * torch.log(hard_probs + 1e-7), dim=-1) |
| | ).sum() |
| |
|
| | if produce_targets: |
| | result["targets"] = idx |
| |
|
| | if self.time_first: |
| | x = x.transpose(1, 2) |
| | result["x"] = x |
| |
|
| | ze = ze.float() |
| | zq = zq.float() |
| | latent_loss = self.mse_mean(zq, ze.detach()) |
| | commitment_loss = self.mse_mean(ze, zq.detach()) |
| |
|
| | result["kmeans_loss"] = latent_loss + self.gamma * commitment_loss |
| |
|
| | return result |
| |
|
