From 001b34f9a295fa8e828a30b1f820c4ef6eef0d32 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Wed, 17 Nov 2021 22:40:26 +0100 Subject: Update vq transformer lit model --- text_recognizer/models/vq_transformer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) (limited to 'text_recognizer') diff --git a/text_recognizer/models/vq_transformer.py b/text_recognizer/models/vq_transformer.py index c07efa4..8ec28fd 100644 --- a/text_recognizer/models/vq_transformer.py +++ b/text_recognizer/models/vq_transformer.py @@ -12,8 +12,6 @@ from text_recognizer.models.transformer import TransformerLitModel class VqTransformerLitModel(TransformerLitModel): """A PyTorch Lightning model for transformer networks.""" - alpha: float = attr.ib(default=1.0) - def forward(self, data: Tensor) -> Tensor: """Forward pass with the transformer network.""" return self.predict(data) @@ -22,7 +20,7 @@ class VqTransformerLitModel(TransformerLitModel): """Training step.""" data, targets = batch logits, commitment_loss = self.network(data, targets[:, :-1]) - loss = self.loss_fn(logits, targets[:, 1:]) + self.alpha * commitment_loss + loss = self.loss_fn(logits, targets[:, 1:]) + commitment_loss self.log("train/loss", loss) self.log("train/commitment_loss", commitment_loss) return loss @@ -31,7 +29,7 @@ class VqTransformerLitModel(TransformerLitModel): """Validation step.""" data, targets = batch logits, commitment_loss = self.network(data, targets[:, :-1]) - loss = self.loss_fn(logits, targets[:, 1:]) + self.alpha * commitment_loss + loss = self.loss_fn(logits, targets[:, 1:]) + commitment_loss self.log("val/loss", loss, prog_bar=True) self.log("val/commitment_loss", commitment_loss) -- cgit v1.2.3-70-g09d2