diff options
author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2023-09-03 01:14:41 +0200 |
---|---|---|
committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2023-09-03 01:14:41 +0200 |
commit | 53cfc21cffa4e877ad0959170b47b690d2fdb40f (patch) | |
tree | 88eaabb75a5ca6fba3c43655ccb754a9f0c5d649 | |
parent | d020059f2f71fe7c25765dde9d535195c09ece01 (diff) |
Update lit models
-rw-r--r-- | text_recognizer/model/base.py | 7 | ||||
-rw-r--r-- | text_recognizer/model/greedy_decoder.py | 2 | ||||
-rw-r--r-- | text_recognizer/model/transformer.py | 6 |
3 files changed, 4 insertions, 11 deletions
diff --git a/text_recognizer/model/base.py b/text_recognizer/model/base.py index adcb8da..9a751bf 100644 --- a/text_recognizer/model/base.py +++ b/text_recognizer/model/base.py @@ -7,7 +7,6 @@ from loguru import logger as log from omegaconf import DictConfig import pytorch_lightning as L from torch import nn, Tensor -from torchmetrics import Accuracy from text_recognizer.data.tokenizer import Tokenizer @@ -24,17 +23,11 @@ class LitBase(L.LightningModule): tokenizer: Tokenizer, ) -> None: super().__init__() - self.network = network self.loss_fn = loss_fn self.optimizer_config = optimizer_config self.lr_scheduler_config = lr_scheduler_config self.tokenizer = tokenizer - ignore_index = int(self.tokenizer.get_value("<p>")) - # Placeholders - self.train_acc = Accuracy(mdmc_reduce="samplewise", ignore_index=ignore_index) - self.val_acc = Accuracy(mdmc_reduce="samplewise", ignore_index=ignore_index) - self.test_acc = Accuracy(mdmc_reduce="samplewise", ignore_index=ignore_index) def optimizer_zero_grad( self, diff --git a/text_recognizer/model/greedy_decoder.py b/text_recognizer/model/greedy_decoder.py index 2c4c16e..8d55a02 100644 --- a/text_recognizer/model/greedy_decoder.py +++ b/text_recognizer/model/greedy_decoder.py @@ -34,7 +34,7 @@ class GreedyDecoder: for i in range(1, self.max_output_len): tokens = indecies[:, :i] # (B, Sy) logits = self.network.decode(tokens, img_features) # [ B, N, C ] - indecies_ = torch.argmax(logits, dim=2) # [ B, N ] + indecies_ = logits.argmax(dim=2) # [ B, N ] indecies[:, i] = indecies_[:, -1] # Early stopping of prediction loop if token is end or padding token. diff --git a/text_recognizer/model/transformer.py b/text_recognizer/model/transformer.py index ae6947c..598d995 100644 --- a/text_recognizer/model/transformer.py +++ b/text_recognizer/model/transformer.py @@ -8,7 +8,7 @@ from torchmetrics import CharErrorRate, WordErrorRate from .greedy_decoder import GreedyDecoder from text_recognizer.data.tokenizer import Tokenizer -from text_recognizer.model.base import LitBase +from .base import LitBase class LitTransformer(LitBase): @@ -45,7 +45,7 @@ class LitTransformer(LitBase): logits = self.network(data, targets) # [B, N, C] return logits.permute(0, 2, 1) # [B, C, N] - def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> Tensor: + def training_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> dict: """Training step.""" data, targets = batch logits = self.teacher_forward(data, targets[:, :-1]) @@ -61,7 +61,7 @@ class LitTransformer(LitBase): ), self.tokenizer.batch_decode(targets) outputs.update({"predictions": preds, "ground_truths": gts}) - return loss + return outputs def validation_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> dict: """Validation step.""" |