summaryrefslogtreecommitdiff
path: root/text_recognizer/models
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer/models')
-rw-r--r--text_recognizer/models/base.py10
-rw-r--r--text_recognizer/models/transformer.py8
2 files changed, 9 insertions, 9 deletions
diff --git a/text_recognizer/models/base.py b/text_recognizer/models/base.py
index f4df3f3..3c1919e 100644
--- a/text_recognizer/models/base.py
+++ b/text_recognizer/models/base.py
@@ -2,7 +2,7 @@
from typing import Any, Dict, List, Union, Tuple, Type
import madgrad
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, OmegaConf
import pytorch_lightning as pl
import torch
from torch import nn
@@ -16,9 +16,9 @@ class LitBaseModel(pl.LightningModule):
def __init__(
self,
network: Type[nn.Module],
- optimizer: Union[OmegaConf, Dict],
- lr_scheduler: Union[OmegaConf, Dict],
- criterion: Union[OmegaConf, Dict],
+ optimizer: Union[DictConfig, Dict],
+ lr_scheduler: Union[DictConfig, Dict],
+ criterion: Union[DictConfig, Dict],
monitor: str = "val_loss",
) -> None:
super().__init__()
@@ -34,7 +34,7 @@ class LitBaseModel(pl.LightningModule):
self.test_acc = torchmetrics.Accuracy()
@staticmethod
- def configure_criterion(criterion: Union[OmegaConf, Dict]) -> Type[nn.Module]:
+ def configure_criterion(criterion: Union[DictConfig, Dict]) -> Type[nn.Module]:
"""Returns a loss functions."""
criterion = OmegaConf.create(criterion)
args = {} or criterion.args
diff --git a/text_recognizer/models/transformer.py b/text_recognizer/models/transformer.py
index 983e274..b23685b 100644
--- a/text_recognizer/models/transformer.py
+++ b/text_recognizer/models/transformer.py
@@ -1,7 +1,7 @@
"""PyTorch Lightning model for base Transformers."""
from typing import Dict, List, Optional, Union, Tuple
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, OmegaConf
import pytorch_lightning as pl
import torch
from torch import nn
@@ -20,9 +20,9 @@ class LitTransformerModel(LitBaseModel):
def __init__(
self,
network: Type[nn, Module],
- optimizer: Union[OmegaConf, Dict],
- lr_scheduler: Union[OmegaConf, Dict],
- criterion: Union[OmegaConf, Dict],
+ optimizer: Union[DictConfig, Dict],
+ lr_scheduler: Union[DictConfig, Dict],
+ criterion: Union[DictConfig, Dict],
monitor: str = "val_loss",
mapping: Optional[List[str]] = None,
) -> None: