Update metrics

6 files changed, 71 insertions, 48 deletions

diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index f917635..bb4e695 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -6,7 +6,7 @@ import torch
 from loguru import logger as log
 from omegaconf import DictConfig
 from pytorch_lightning import LightningModule
-from torch import Tensor, nn
+from torch import nn, Tensor
 from torchmetrics import Accuracy
 
 from text_recognizer.data.mappings import EmnistMapping
@@ -22,6 +22,7 @@ class LitBase(LightningModule):
         optimizer_config: DictConfig,
         lr_scheduler_config: Optional[DictConfig],
         mapping: EmnistMapping,
+        ignore_index: Optional[int] = None,
     ) -> None:
         super().__init__()
 
@@ -32,9 +33,9 @@ class LitBase(LightningModule):
         self.mapping = mapping
 
         # Placeholders
-        self.train_acc = Accuracy(mdmc_reduce="samplewise")
-        self.val_acc = Accuracy(mdmc_reduce="samplewise")
-        self.test_acc = Accuracy(mdmc_reduce="samplewise")
+        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/models/metrics.py b/text_recognizer/models/metrics.py
deleted file mode 100644
index 3cb16b5..0000000
--- a/text_recognizer/models/metrics.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""Character Error Rate (CER)."""
-from typing import Sequence
-
-import editdistance
-import torch
-from torch import Tensor
-from torchmetrics import Metric
-
-
-class CharacterErrorRate(Metric):
-    """Character error rate metric, computed using Levenshtein distance."""
-
-    def __init__(self, ignore_indices: Sequence[Tensor]) -> None:
-        super().__init__()
-        self.ignore_indices = set(ignore_indices)
-        self.add_state("error", default=torch.tensor(0.0), dist_reduce_fx="sum")
-        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")
-        self.error: Tensor
-        self.total: Tensor
-
-    def update(self, preds: Tensor, targets: Tensor) -> None:
-        """Update CER."""
-        bsz = preds.shape[0]
-        for index in range(bsz):
-            pred = [p for p in preds[index].tolist() if p not in self.ignore_indices]
-            target = [
-                t for t in targets[index].tolist() if t not in self.ignore_indices
-            ]
-            distance = editdistance.distance(pred, target)
-            error = distance / max(len(pred), len(target))
-            self.error += error
-        self.total += bsz
-
-    def compute(self) -> Tensor:
-        """Compute CER."""
-        return self.error / self.total
diff --git a/text_recognizer/models/metrics/__init__.py b/text_recognizer/models/metrics/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/text_recognizer/models/metrics/__init__.py
diff --git a/text_recognizer/models/metrics/cer.py b/text_recognizer/models/metrics/cer.py
new file mode 100644
index 0000000..238ecc3
--- /dev/null
+++ b/text_recognizer/models/metrics/cer.py
@@ -0,0 +1,23 @@
+"""Character Error Rate (CER)."""
+from typing import Sequence
+
+import torch
+import torchmetrics
+
+
+class CharacterErrorRate(torchmetrics.CharErrorRate):
+    """Character error rate metric, allowing for tokens to be ignored."""
+
+    def __init__(self, ignore_tokens: Sequence[int], *args):
+        super().__init__(*args)
+        self.ignore_tokens = set(ignore_tokens)
+
+    def update(self, preds: torch.Tensor, targets: torch.Tensor) -> None:
+        preds_l = [
+            [t for t in pred if t not in self.ignore_tokens] for pred in preds.tolist()
+        ]
+        targets_l = [
+            [t for t in target if t not in self.ignore_tokens]
+            for target in targets.tolist()
+        ]
+        super().update(preds_l, targets_l)
diff --git a/text_recognizer/models/metrics/wer.py b/text_recognizer/models/metrics/wer.py
new file mode 100644
index 0000000..78f5854
--- /dev/null
+++ b/text_recognizer/models/metrics/wer.py
@@ -0,0 +1,23 @@
+"""Character Error Rate (CER)."""
+from typing import Sequence
+
+import torch
+import torchmetrics
+
+
+class WordErrorRate(torchmetrics.WordErrorRate):
+    """Character error rate metric, allowing for tokens to be ignored."""
+
+    def __init__(self, ignore_tokens: Sequence[int], *args):
+        super().__init__(*args)
+        self.ignore_tokens = set(ignore_tokens)
+
+    def update(self, preds: torch.Tensor, targets: torch.Tensor) -> None:
+        preds_l = [
+            [t for t in pred if t not in self.ignore_tokens] for pred in preds.tolist()
+        ]
+        targets_l = [
+            [t for t in target if t not in self.ignore_tokens]
+            for target in targets.tolist()
+        ]
+        super().update(preds_l, targets_l)
diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index dcec756..2c74b7e 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -3,11 +3,12 @@ from typing import Optional, Tuple, Type
 
 import torch
 from omegaconf import DictConfig
-from torch import Tensor, nn
+from torch import nn, Tensor
 
 from text_recognizer.data.mappings import EmnistMapping
 from text_recognizer.models.base import LitBase
-from text_recognizer.models.metrics import CharacterErrorRate
+from text_recognizer.models.metrics.cer import CharacterErrorRate
+from text_recognizer.models.metrics.wer import WordErrorRate
 
 
 class LitTransformer(LitBase):
@@ -18,16 +19,13 @@ class LitTransformer(LitBase):
         network: Type[nn.Module],
         loss_fn: Type[nn.Module],
         optimizer_config: DictConfig,
-        lr_scheduler_config: Optional[DictConfig],
         mapping: EmnistMapping,
+        lr_scheduler_config: Optional[DictConfig] = None,
         max_output_len: int = 682,
         start_token: str = "<s>",
         end_token: str = "<e>",
         pad_token: str = "<p>",
     ) -> None:
-        super().__init__(
-            network, loss_fn, optimizer_config, lr_scheduler_config, mapping
-        )
         self.max_output_len = max_output_len
         self.start_token = start_token
         self.end_token = end_token
@@ -38,6 +36,16 @@ class LitTransformer(LitBase):
         self.ignore_indices = set([self.start_index, self.end_index, self.pad_index])
         self.val_cer = CharacterErrorRate(self.ignore_indices)
         self.test_cer = CharacterErrorRate(self.ignore_indices)
+        self.val_wer = WordErrorRate(self.ignore_indices)
+        self.test_wer = WordErrorRate(self.ignore_indices)
+        super().__init__(
+            network,
+            loss_fn,
+            optimizer_config,
+            lr_scheduler_config,
+            mapping,
+            self.pad_index,
+        )
 
     def forward(self, data: Tensor) -> Tensor:
         """Forward pass with the transformer network."""
@@ -59,6 +67,8 @@ class LitTransformer(LitBase):
         self.log("val/acc", self.val_acc, on_step=False, on_epoch=True)
         self.val_cer(preds, targets)
         self.log("val/cer", self.val_cer, on_step=False, on_epoch=True, prog_bar=True)
+        self.val_wer(preds, targets)
+        self.log("val/wer", self.val_wer, on_step=False, on_epoch=True, prog_bar=True)
 
     def test_step(self, batch: Tuple[Tensor, Tensor], batch_idx: int) -> None:
         """Test step."""
@@ -66,10 +76,12 @@ class LitTransformer(LitBase):
 
         # Compute the text prediction.
         pred = self(data)
-        self.test_cer(pred, targets)
-        self.log("test/cer", self.test_cer, on_step=False, on_epoch=True, prog_bar=True)
         self.test_acc(pred, targets)
         self.log("test/acc", self.test_acc, on_step=False, on_epoch=True)
+        self.test_cer(pred, targets)
+        self.log("test/cer", self.test_cer, on_step=False, on_epoch=True, prog_bar=True)
+        self.test_wer(pred, targets)
+        self.log("test/wer", self.test_wer, on_step=False, on_epoch=True, prog_bar=True)
 
     @torch.no_grad()
     def predict(self, x: Tensor) -> Tensor: