summaryrefslogtreecommitdiff
path: root/text_recognizer/models
diff options
context:
space:
mode:
authorGustaf Rydholm <gustaf.rydholm@gmail.com>2022-06-05 23:39:11 +0200
committerGustaf Rydholm <gustaf.rydholm@gmail.com>2022-06-05 23:39:11 +0200
commit65df6a72b002c4b23d6f2eb545839e157f7f2aa0 (patch)
treed78df1d7143dc9ff9e29afd4fd6bc7490bc79418 /text_recognizer/models
parent8bc4b4cab00a2777a748c10fca9b3ee01e32277c (diff)
Remove attrs
Diffstat (limited to 'text_recognizer/models')
-rw-r--r--text_recognizer/models/base.py36
-rw-r--r--text_recognizer/models/metrics.py13
-rw-r--r--text_recognizer/models/transformer.py29
3 files changed, 36 insertions, 42 deletions
diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index bf3bc08..63fe5a7 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -1,7 +1,6 @@
"""Base PyTorch Lightning model."""
from typing import Any, Dict, List, Optional, Tuple, Type
-from attrs import define, field
import hydra
from loguru import logger as log
from omegaconf import DictConfig
@@ -9,31 +8,34 @@ from pytorch_lightning import LightningModule
import torch
from torch import nn
from torch import Tensor
-import torchmetrics
+from torchmetrics import Accuracy
from text_recognizer.data.mappings.base import AbstractMapping
-@define(eq=False)
class BaseLitModel(LightningModule):
"""Abstract PyTorch Lightning class."""
- def __attrs_pre_init__(self) -> None:
- """Pre init constructor."""
+ def __init__(
+ self,
+ network: Type[nn.Module],
+ loss_fn: Type[nn.Module],
+ optimizer_configs: DictConfig,
+ lr_scheduler_configs: Optional[DictConfig],
+ mapping: Type[AbstractMapping],
+ ) -> None:
super().__init__()
- network: Type[nn.Module] = field()
- loss_fn: Type[nn.Module] = field()
- optimizer_configs: DictConfig = field()
- lr_scheduler_configs: Optional[DictConfig] = field()
- mapping: Type[AbstractMapping] = field()
-
- # Placeholders
- train_acc: torchmetrics.Accuracy = field(
- init=False, default=torchmetrics.Accuracy()
- )
- val_acc: torchmetrics.Accuracy = field(init=False, default=torchmetrics.Accuracy())
- test_acc: torchmetrics.Accuracy = field(init=False, default=torchmetrics.Accuracy())
+ self.network = network
+ self.loss_fn = loss_fn
+ self.optimizer_configs = optimizer_configs
+ self.lr_scheduler_configs = lr_scheduler_configs
+ self.mapping = mapping
+
+ # Placeholders
+ self.train_acc = Accuracy()
+ self.val_acc = Accuracy()
+ self.test_acc = Accuracy()
def optimizer_zero_grad(
self,
diff --git a/text_recognizer/models/metrics.py b/text_recognizer/models/metrics.py
index e59a830..3cb16b5 100644
--- a/text_recognizer/models/metrics.py
+++ b/text_recognizer/models/metrics.py
@@ -1,25 +1,22 @@
"""Character Error Rate (CER)."""
-from typing import Set
+from typing import Sequence
-from attrs import define, field
import editdistance
import torch
from torch import Tensor
from torchmetrics import Metric
-@define(eq=False)
class CharacterErrorRate(Metric):
"""Character error rate metric, computed using Levenshtein distance."""
- ignore_indices: Set[Tensor] = field(converter=set)
- error: Tensor = field(init=False)
- total: Tensor = field(init=False)
-
- def __attrs_post_init__(self) -> None:
+ 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."""
diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index c5120fe..9537dd9 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -1,7 +1,6 @@
"""PyTorch Lightning model for base Transformers."""
from typing import Set, Tuple
-from attrs import define, field
import torch
from torch import Tensor
@@ -9,25 +8,21 @@ from text_recognizer.models.base import BaseLitModel
from text_recognizer.models.metrics import CharacterErrorRate
-@define(auto_attribs=True, eq=False)
class TransformerLitModel(BaseLitModel):
"""A PyTorch Lightning model for transformer networks."""
- max_output_len: int = field(default=451)
- start_token: str = field(default="<s>")
- end_token: str = field(default="<e>")
- pad_token: str = field(default="<p>")
-
- start_index: int = field(init=False)
- end_index: int = field(init=False)
- pad_index: int = field(init=False)
-
- ignore_indices: Set[Tensor] = field(init=False)
- val_cer: CharacterErrorRate = field(init=False)
- test_cer: CharacterErrorRate = field(init=False)
-
- def __attrs_post_init__(self) -> None:
- """Post init configuration."""
+ def __init__(
+ self,
+ max_output_len: int = 451,
+ start_token: str = "<s>",
+ end_token: str = "<e>",
+ pad_token: str = "<p>",
+ ) -> None:
+ super().__init__()
+ self.max_output_len = max_output_len
+ self.start_token = start_token
+ self.end_token = end_token
+ self.pad_token = pad_token
self.start_index = int(self.mapping.get_index(self.start_token))
self.end_index = int(self.mapping.get_index(self.end_token))
self.pad_index = int(self.mapping.get_index(self.pad_token))