| import numpy as np |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from .layers import MLPLayers |
| from .rq import ResidualVectorQuantizer |
|
|
|
|
| class RQVAE(nn.Module): |
| def __init__(self, |
| in_dim=768, |
| |
| num_emb_list=None, |
| e_dim=64, |
| |
| layers=None, |
| dropout_prob=0.0, |
| bn=False, |
| loss_type="mse", |
| quant_loss_weight=1.0, |
| kmeans_init=False, |
| kmeans_iters=100, |
| |
| sk_epsilons=None, |
| sk_iters=100, |
| ): |
| super(RQVAE, self).__init__() |
|
|
| self.in_dim = in_dim |
| self.num_emb_list = num_emb_list |
| self.e_dim = e_dim |
|
|
| self.layers = layers |
| self.dropout_prob = dropout_prob |
| self.bn = bn |
| self.loss_type = loss_type |
| self.quant_loss_weight=quant_loss_weight |
| self.kmeans_init = kmeans_init |
| self.kmeans_iters = kmeans_iters |
| self.sk_epsilons = sk_epsilons |
| self.sk_iters = sk_iters |
|
|
| self.encode_layer_dims = [self.in_dim] + self.layers + [self.e_dim] |
| self.encoder = MLPLayers(layers=self.encode_layer_dims, |
| dropout=self.dropout_prob,bn=self.bn) |
|
|
| self.rq = ResidualVectorQuantizer(num_emb_list, e_dim, |
| kmeans_init = self.kmeans_init, |
| kmeans_iters = self.kmeans_iters, |
| sk_epsilons=self.sk_epsilons, |
| sk_iters=self.sk_iters,) |
|
|
| self.decode_layer_dims = self.encode_layer_dims[::-1] |
| self.decoder = MLPLayers(layers=self.decode_layer_dims, |
| dropout=self.dropout_prob,bn=self.bn) |
|
|
| def forward(self, x, use_sk=True): |
| |
| x = self.encoder(x) |
| x_q, rq_loss, indices = self.rq(x,use_sk=use_sk) |
| out = self.decoder(x_q) |
| |
|
|
| return out, rq_loss, indices |
|
|
| @torch.no_grad() |
| def get_indices(self, xs, use_sk=False): |
| x_e = self.encoder(xs) |
| _, _, indices = self.rq(x_e, use_sk=use_sk) |
| return indices |
|
|
| def compute_loss(self, out, quant_loss, xs=None): |
|
|
| if self.loss_type == 'mse': |
| loss_recon = F.mse_loss(out, xs, reduction='mean') |
| elif self.loss_type == 'l1': |
| loss_recon = F.l1_loss(out, xs, reduction='mean') |
| else: |
| raise ValueError('incompatible loss type') |
|
|
| loss_total = loss_recon + self.quant_loss_weight * quant_loss |
|
|
| return loss_total, loss_recon |
