summaryrefslogtreecommitdiff
path: root/text_recognizer/models/vqgan.py
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer/models/vqgan.py')
-rw-r--r--text_recognizer/models/vqgan.py24
1 files changed, 2 insertions, 22 deletions
diff --git a/text_recognizer/models/vqgan.py b/text_recognizer/models/vqgan.py
index 8ff65cc..80653b6 100644
--- a/text_recognizer/models/vqgan.py
+++ b/text_recognizer/models/vqgan.py
@@ -9,7 +9,7 @@ from text_recognizer.criterions.vqgan_loss import VQGANLoss
@attr.s(auto_attribs=True, eq=False)
-class VQVAELitModel(BaseLitModel):
+class VQGANLitModel(BaseLitModel):
"""A PyTorch Lightning model for transformer networks."""
loss_fn: VQGANLoss = attr.ib()
@@ -26,7 +26,6 @@ class VQVAELitModel(BaseLitModel):
data, _ = batch
reconstructions, vq_loss = self(data)
- loss = self.loss_fn(reconstructions, data)
if optimizer_idx == 0:
loss, log = self.loss_fn(
@@ -81,14 +80,6 @@ class VQVAELitModel(BaseLitModel):
self.log(
"val/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True
)
- self.log(
- "val/rec_loss",
- log["val/rec_loss"],
- prog_bar=True,
- logger=True,
- on_step=True,
- on_epoch=True,
- )
self.log_dict(log)
_, log = self.loss_fn(
@@ -105,24 +96,13 @@ class VQVAELitModel(BaseLitModel):
data, _ = batch
reconstructions, vq_loss = self(data)
- loss, log = self.loss_fn(
+ _, log = self.loss_fn(
data=data,
reconstructions=reconstructions,
vq_loss=vq_loss,
optimizer_idx=0,
stage="test",
)
- self.log(
- "test/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True
- )
- self.log(
- "test/rec_loss",
- log["test/rec_loss"],
- prog_bar=True,
- logger=True,
- on_step=True,
- on_epoch=True,
- )
self.log_dict(log)
_, log = self.loss_fn(