summaryrefslogtreecommitdiff
path: root/text_recognizer/models
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2021-08-04 05:03:51 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2021-08-04 05:03:51 +0200
commitd3afa310f77f47553586eeee58e3d3345a754e2c (patch)
tree08b7de1daf2550852d0a1e4d4d75202f14bb03d4 /text_recognizer/models
parent65d5f6c694e73792e40ed693a1381a792da8d277 (diff)
New VQVAE
Diffstat (limited to 'text_recognizer/models')
-rw-r--r--text_recognizer/models/vqvae.py16
1 files changed, 9 insertions, 7 deletions
diff --git a/text_recognizer/models/vqvae.py b/text_recognizer/models/vqvae.py
index 22da018..5890fd9 100644
--- a/text_recognizer/models/vqvae.py
+++ b/text_recognizer/models/vqvae.py
@@ -14,31 +14,33 @@ from text_recognizer.models.base import BaseLitModel
class VQVAELitModel(BaseLitModel):
"""A PyTorch Lightning model for transformer networks."""
+ latent_loss_weight: float = attr.ib(default=0.25)
+
def forward(self, data: Tensor) -> Tensor:
"""Forward pass with the transformer network."""
- return self.network.predict(data)
+ return self.network(data)
def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Tensor:
"""Training step."""
data, _ = batch
- reconstructions, vq_loss = self.network(data)
+ reconstructions, vq_loss = self(data)
loss = self.loss_fn(reconstructions, data)
- loss += vq_loss
+ loss += self.latent_loss_weight * vq_loss
self.log("train/loss", loss)
return loss
def validation_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None:
"""Validation step."""
data, _ = batch
- reconstructions, vq_loss = self.network(data)
+ reconstructions, vq_loss = self(data)
loss = self.loss_fn(reconstructions, data)
- loss += vq_loss
+ loss += self.latent_loss_weight * vq_loss
self.log("val/loss", loss, prog_bar=True)
def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None:
"""Test step."""
data, _ = batch
- reconstructions, vq_loss = self.network(data)
+ reconstructions, vq_loss = self(data)
loss = self.loss_fn(reconstructions, data)
- loss += vq_loss
+ loss += self.latent_loss_weight * vq_loss
self.log("test/loss", loss)