summaryrefslogtreecommitdiff
path: root/src/training
diff options
context:
space:
mode:
Diffstat (limited to 'src/training')
-rw-r--r--src/training/run_experiment.py8
-rw-r--r--src/training/trainer/train.py4
2 files changed, 8 insertions, 4 deletions
diff --git a/src/training/run_experiment.py b/src/training/run_experiment.py
index 55a9572..a883b45 100644
--- a/src/training/run_experiment.py
+++ b/src/training/run_experiment.py
@@ -21,8 +21,9 @@ from training.trainer.train import Trainer
import wandb
import yaml
-
+import text_recognizer.models
from text_recognizer.models import Model
+import text_recognizer.networks
from text_recognizer.networks.loss import loss as custom_loss_module
EXPERIMENTS_DIRNAME = Path(__file__).parents[0].resolve() / "experiments"
@@ -77,13 +78,12 @@ def _load_modules_and_arguments(experiment_config: Dict,) -> Tuple[Callable, Dic
dataset_ = dataset_args["type"]
# Import the model module and model arguments.
- models_module = importlib.import_module("text_recognizer.models")
- model_class_ = getattr(models_module, experiment_config["model"])
+ model_class_ = getattr(text_recognizer.models, experiment_config["model"])
# Import metrics.
metric_fns_ = (
{
- metric: getattr(models_module, metric)
+ metric: getattr(text_recognizer.networks, metric)
for metric in experiment_config["metrics"]
}
if experiment_config["metrics"] is not None
diff --git a/src/training/trainer/train.py b/src/training/trainer/train.py
index 223d9c6..8ae994a 100644
--- a/src/training/trainer/train.py
+++ b/src/training/trainer/train.py
@@ -3,6 +3,7 @@
from pathlib import Path
import time
from typing import Dict, List, Optional, Tuple, Type
+import warnings
from einops import rearrange
from loguru import logger
@@ -23,6 +24,9 @@ torch.manual_seed(4711)
torch.cuda.manual_seed(4711)
+warnings.filterwarnings("ignore")
+
+
class Trainer:
"""Trainer for training PyTorch models."""