From 5a78fc2e33c28968a69d033cb10d638f4f63fed1 Mon Sep 17 00:00:00 2001
From: aktersnurra <gustaf.rydholm@gmail.com>
Date: Sun, 5 Jul 2020 22:27:08 +0200
Subject: Working on getting experiment loop.

---
 src/text_recognizer/models/base.py            | 12 ++++++------
 src/text_recognizer/models/character_model.py |  8 ++++++--
 src/text_recognizer/models/metrics.py         | 19 +++++++++++++++++++
 src/text_recognizer/models/util.py            | 19 -------------------
 4 files changed, 31 insertions(+), 27 deletions(-)
 create mode 100644 src/text_recognizer/models/metrics.py
 delete mode 100644 src/text_recognizer/models/util.py

(limited to 'src/text_recognizer/models')

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:
diff --git a/src/text_recognizer/models/character_model.py b/src/text_recognizer/models/character_model.py
index 1570344..fd69bf2 100644
--- a/src/text_recognizer/models/character_model.py
+++ b/src/text_recognizer/models/character_model.py
@@ -32,17 +32,21 @@ class CharacterModel(Model):
 
         super().__init__(
             network_fn,
-            data_loader_args,
             network_args,
+            data_loader_args,
             metrics,
             criterion,
+            criterion_args,
             optimizer,
+            optimizer_args,
+            lr_scheduler,
+            lr_scheduler_args,
             device,
         )
         self.emnist_mapping = self.mapping()
         self.eval()
 
-    def mapping(self) -> Dict:
+    def mapping(self) -> Dict[int, str]:
         """Mapping between integers and classes."""
         mapping = load_emnist_mapping()
         return mapping
diff --git a/src/text_recognizer/models/metrics.py b/src/text_recognizer/models/metrics.py
new file mode 100644
index 0000000..e2a30a9
--- /dev/null
+++ b/src/text_recognizer/models/metrics.py
@@ -0,0 +1,19 @@
+"""Utility functions for models."""
+
+import torch
+
+
+def accuracy(outputs: torch.Tensor, labels: torch.Tensro) -> float:
+    """Computes the accuracy.
+
+    Args:
+        outputs (torch.Tensor): The output from the network.
+        labels (torch.Tensor): Ground truth labels.
+
+    Returns:
+        float: The accuracy for the batch.
+
+    """
+    _, predicted = torch.max(outputs.data, dim=1)
+    acc = (predicted == labels).sum().item() / labels.shape[0]
+    return acc
diff --git a/src/text_recognizer/models/util.py b/src/text_recognizer/models/util.py
deleted file mode 100644
index 905fe7b..0000000
--- a/src/text_recognizer/models/util.py
+++ /dev/null
@@ -1,19 +0,0 @@
-"""Utility functions for models."""
-
-import torch
-
-
-def accuracy(outputs: torch.Tensor, labels: torch.Tensro) -> float:
-    """Short summary.
-
-    Args:
-        outputs (torch.Tensor): The output from the network.
-        labels (torch.Tensor): Ground truth labels.
-
-    Returns:
-        float: The accuracy for the batch.
-
-    """
-    _, predicted = torch.max(outputs.data, dim=1)
-    acc = (predicted == labels).sum().item() / labels.shape[0]
-    return acc
-- 
cgit v1.2.3-70-g09d2