| | 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 |
