From 5a78fc2e33c28968a69d033cb10d638f4f63fed1 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Sun, 5 Jul 2020 22:27:08 +0200 Subject: Working on getting experiment loop. --- src/text_recognizer/models/base.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'src/text_recognizer/models/base.py') diff --git a/src/text_recognizer/models/base.py b/src/text_recognizer/models/base.py index 736af7b..0cc531a 100644 --- a/src/text_recognizer/models/base.py +++ b/src/text_recognizer/models/base.py @@ -10,7 +10,6 @@ import torch from torch import nn from torchsummary import summary -from text_recognizer.dataset.data_loader import fetch_data_loader WEIGHT_DIRNAME = Path(__file__).parents[1].resolve() / "weights" @@ -22,6 +21,7 @@ class Model(ABC): self, network_fn: Callable, network_args: Dict, + data_loader: Optional[Callable] = None, data_loader_args: Optional[Dict] = None, metrics: Optional[Dict] = None, criterion: Optional[Callable] = None, @@ -32,12 +32,13 @@ class Model(ABC): lr_scheduler_args: Optional[Dict] = None, device: Optional[str] = None, ) -> None: - """Base class, to be inherited by predictors for specific type of data. + """Base class, to be inherited by model for specific type of data. Args: network_fn (Callable): The PyTorch network. network_args (Dict): Arguments for the network. - data_loader_args (Optional[Dict]): Arguments for the data loader. + data_loader (Optional[Callable]): A function that fetches train and val DataLoader. + data_loader_args (Optional[Dict]): Arguments for the DataLoader. metrics (Optional[Dict]): Metrics to evaluate the performance with. Defaults to None. criterion (Optional[Callable]): The criterion to evaulate the preformance of the network. Defaults to None. @@ -53,8 +54,8 @@ class Model(ABC): # Fetch data loaders. if data_loader_args is not None: - self._data_loaders = fetch_data_loader(**data_loader_args) - dataset_name = self._data_loaders.items()[0].dataset.__name__ + self._data_loaders = data_loader(**data_loader_args) + dataset_name = self._data_loaders.__name__ else: dataset_name = "" self._data_loaders = None @@ -210,7 +211,6 @@ class Model(ABC): logger.debug( f"Found a new best {val_metric}. Saving best checkpoint and weights." ) - self.save_weights() shutil.copyfile(filepath, str(path / "best.pt")) def load_weights(self) -> None: -- cgit v1.2.3-70-g09d2