summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/vqvae/vqvae.py
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2021-11-17 22:43:12 +0100
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2021-11-17 22:43:12 +0100
commit700ce6ed83867601de0ae55032afdd5e12438258 (patch)
treecbf2e47a586a0a84b4fd69fbc061cd6de7f06904 /text_recognizer/networks/vqvae/vqvae.py
parent2417288c9fe96264da708ce8d13ac7bc2faf83e3 (diff)
Update vqvae with new quantizer
Diffstat (limited to 'text_recognizer/networks/vqvae/vqvae.py')
-rw-r--r--text_recognizer/networks/vqvae/vqvae.py26
1 files changed, 6 insertions, 20 deletions
diff --git a/text_recognizer/networks/vqvae/vqvae.py b/text_recognizer/networks/vqvae/vqvae.py
index d876ca1..5560e12 100644
--- a/text_recognizer/networks/vqvae/vqvae.py
+++ b/text_recognizer/networks/vqvae/vqvae.py
@@ -4,7 +4,7 @@ from typing import Tuple
from torch import nn
from torch import Tensor
-from text_recognizer.networks.vqvae.quantizer import VectorQuantizer
+from text_recognizer.networks.quantizer.quantizer import VectorQuantizer
class VQVAE(nn.Module):
@@ -14,39 +14,25 @@ class VQVAE(nn.Module):
self,
encoder: nn.Module,
decoder: nn.Module,
- hidden_dim: int,
- embedding_dim: int,
- num_embeddings: int,
- decay: float = 0.99,
+ quantizer: VectorQuantizer,
) -> None:
super().__init__()
self.encoder = encoder
self.decoder = decoder
- self.pre_codebook_conv = nn.Conv2d(
- in_channels=hidden_dim, out_channels=embedding_dim, kernel_size=1
- )
- self.post_codebook_conv = nn.Conv2d(
- in_channels=embedding_dim, out_channels=hidden_dim, kernel_size=1
- )
- self.quantizer = VectorQuantizer(
- num_embeddings=num_embeddings, embedding_dim=embedding_dim, decay=decay,
- )
+ self.quantizer = quantizer
def encode(self, x: Tensor) -> Tensor:
"""Encodes input to a latent code."""
- z_e = self.encoder(x)
- return self.pre_codebook_conv(z_e)
+ return self.encoder(x)
def quantize(self, z_e: Tensor) -> Tuple[Tensor, Tensor]:
"""Quantizes the encoded latent vectors."""
- z_q, commitment_loss = self.quantizer(z_e)
+ z_q, _, commitment_loss = self.quantizer(z_e)
return z_q, commitment_loss
def decode(self, z_q: Tensor) -> Tensor:
"""Reconstructs input from latent codes."""
- z = self.post_codebook_conv(z_q)
- x_hat = self.decoder(z)
- return x_hat
+ return self.decoder(z_q)
def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
"""Compresses and decompresses input."""