feat: logging of train and eval better aligned (#105)

finetuner/tuner/keras/__init__.py

@@ -12,6 +12,7 @@
 from ..base import BaseTuner
 from ...helper import DocumentArrayLike
 from ..dataset.helper import get_dataset
+from ..logger import LogGenerator
 
 
 class KerasTuner(BaseTuner):
@@ -62,11 +63,12 @@ def _train(self, data, optimizer, description: str):
         losses = []
         metrics = []
 
-        get_desc_str = (
-            lambda: f'Loss={np.mean(losses):.2f} Accuracy={np.mean(metrics):.2f}'
-        )
+        log_generator = LogGenerator('T', losses, metrics)
+
+        with ProgressBar(
+            description, message_on_done=log_generator, final_line_feed=False
+        ) as p:
 
-        with ProgressBar(description, message_on_done=get_desc_str) as p:
             for inputs, label in data:
                 with tf.GradientTape() as tape:
                     outputs = self.wrapped_model(inputs, training=True)
@@ -81,21 +83,19 @@ def _train(self, data, optimizer, description: str):
                 losses.append(loss.numpy())
                 metrics.append(metric.numpy())
 
-                p.update(message=get_desc_str())
+                p.update(message=log_generator())
 
         return losses, metrics
 
-    def _eval(self, data, description: str = 'Evaluating'):
+    def _eval(self, data, description: str = 'Evaluating', train_log: str = ''):
         head_layer = self.head_layer()
 
         losses = []
         metrics = []
 
-        get_desc_str = (
-            lambda: f'Loss={np.mean(losses):.2f} Accuracy={np.mean(metrics):.2f}'
-        )
+        log_generator = LogGenerator('E', losses, metrics, train_log)
 
-        with ProgressBar(description, message_on_done=get_desc_str) as p:
+        with ProgressBar(description, message_on_done=log_generator) as p:
             for inputs, label in data:
                 outputs = self.wrapped_model(inputs, training=False)
                 loss = head_layer.loss_fn(pred_val=outputs, target_val=label)
@@ -104,7 +104,7 @@ def _eval(self, data, description: str = 'Evaluating'):
                 losses.append(loss.numpy())
                 metrics.append(metric.numpy())
 
-                p.update(message=get_desc_str())
+                p.update(message=log_generator())
 
         return losses, metrics
 
@@ -143,7 +143,7 @@ def fit(
             metrics_train.extend(mt)
 
             if eval_data:
-                le, me = self._eval(_eval_data)
+                le, me = self._eval(_eval_data, train_log=LogGenerator('T', lt, mt)())
                 losses_eval.extend(le)
                 metrics_eval.extend(me)
 

finetuner/tuner/logger.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+
+class LogGenerator:
+    def __init__(self, name, losses, metrics, prefix: str = ''):
+        self._losses = losses
+        self._metrics = metrics
+        self._prefix = prefix
+        self._name = name
+
+    def __call__(self):
+        if self._prefix:
+            prefix = f'{self._prefix} | '
+        else:
+            prefix = ''
+        return f'{prefix}{self._name}: {self.get_statistic()}'
+
+    def get_statistic(self):
+        return f'L={self.mean_loss():>8} A={self.mean_metric():>4}'
+
+    def mean_loss(self):
+        return LogGenerator.get_log_value(self._losses)
+
+    def mean_metric(self):
+        return LogGenerator.get_log_value(self._metrics)
+
+    @staticmethod
+    def get_log_value(data):
+        mean = np.mean(data)
+        if mean < 1e5:
+            return f'{mean:.2f}'
+        else:
+            return f'{mean:.2e}'

finetuner/tuner/paddle/__init__.py

@@ -10,6 +10,7 @@
 from ..base import BaseTuner, BaseHead, BaseArityModel
 from ...helper import DocumentArrayLike
 from ..dataset.helper import get_dataset
+from ..logger import LogGenerator
 
 
 class _ArityModel(BaseArityModel, nn.Layer):
@@ -39,17 +40,15 @@ def _get_data_loader(self, inputs, batch_size: int, shuffle: bool):
             shuffle=shuffle,
         )
 
-    def _eval(self, data, description: str = 'Evaluating'):
+    def _eval(self, data, description: str = 'Evaluating', train_log: str = ''):
         self.wrapped_model.eval()
 
         losses = []
         metrics = []
 
-        get_desc_str = (
-            lambda: f'Loss={np.mean(losses):.2f} Accuracy={np.mean(metrics):.2f}'
-        )
+        log_generator = LogGenerator('E', losses, metrics, train_log)
 
-        with ProgressBar(description, message_on_done=get_desc_str) as p:
+        with ProgressBar(description, message_on_done=log_generator) as p:
             for inputs, label in data:
                 outputs = self.wrapped_model(*inputs)
                 loss = self.wrapped_model.loss_fn(outputs, label)
@@ -58,7 +57,7 @@ def _eval(self, data, description: str = 'Evaluating'):
                 losses.append(loss.item())
                 metrics.append(metric.numpy())
 
-                p.update(message=get_desc_str())
+                p.update(message=log_generator())
 
         return losses, metrics
 
@@ -69,11 +68,11 @@ def _train(self, data, optimizer: Optimizer, description: str):
         losses = []
         metrics = []
 
-        get_desc_str = (
-            lambda: f'Loss={np.mean(losses):.2f} Accuracy={float(np.mean(metrics)):.2f}'
-        )
+        log_generator = LogGenerator('T', losses, metrics)
 
-        with ProgressBar(description, message_on_done=get_desc_str) as p:
+        with ProgressBar(
+            description, message_on_done=log_generator, final_line_feed=False
+        ) as p:
             for inputs, label in data:
                 # forward step
                 outputs = self.wrapped_model(*inputs)
@@ -88,7 +87,7 @@ def _train(self, data, optimizer: Optimizer, description: str):
                 losses.append(loss.item())
                 metrics.append(metric.numpy())
 
-                p.update(message=get_desc_str())
+                p.update(message=log_generator())
         return losses, metrics
 
     def fit(
@@ -99,10 +98,8 @@ def fit(
         batch_size: int = 256,
         **kwargs,
     ):
-        model = self.wrapped_model
-
         optimizer = paddle.optimizer.RMSProp(
-            learning_rate=0.01, parameters=model.parameters()
+            learning_rate=0.01, parameters=self.wrapped_model.parameters()
         )
 
         losses_train = []
@@ -127,7 +124,7 @@ def fit(
                     inputs=eval_data, batch_size=batch_size, shuffle=False
                 )
 
-                le, me = self._eval(_data)
+                le, me = self._eval(_data, train_log=LogGenerator('T', lt, mt)())
                 losses_eval.extend(le)
                 metrics_eval.extend(me)
 

finetuner/tuner/pytorch/__init__.py

@@ -1,6 +1,5 @@
 from typing import Optional
 
-import numpy as np
 import torch
 import torch.nn as nn
 from jina.helper import cached_property
@@ -12,6 +11,7 @@
 from ..base import BaseTuner, BaseHead, BaseArityModel
 from ...helper import DocumentArrayLike
 from ..dataset.helper import get_dataset
+from ..logger import LogGenerator
 
 
 class _ArityModel(BaseArityModel, nn.Module):
@@ -41,17 +41,14 @@ def _get_data_loader(self, inputs, batch_size: int, shuffle: bool):
             shuffle=shuffle,
         )
 
-    def _eval(self, data, description: str = 'Evaluating'):
+    def _eval(self, data, description: str = 'Evaluating', train_log: str = ''):
         self.wrapped_model.eval()
 
         losses = []
         metrics = []
+        log_generator = LogGenerator('E', losses, metrics, train_log)
 
-        get_desc_str = (
-            lambda: f'Loss={np.mean(losses):.2f} Accuracy={np.mean(metrics):.2f}'
-        )
-
-        with ProgressBar(description, message_on_done=get_desc_str) as p:
+        with ProgressBar(description, message_on_done=log_generator) as p:
             for inputs, label in data:
                 with torch.inference_mode():
                     outputs = self.wrapped_model(*inputs)
@@ -61,7 +58,7 @@ def _eval(self, data, description: str = 'Evaluating'):
                 losses.append(loss.item())
                 metrics.append(metric.numpy())
 
-                p.update(message=get_desc_str())
+                p.update(message=log_generator())
 
         return losses, metrics
 
@@ -72,11 +69,11 @@ def _train(self, data, optimizer: Optimizer, description: str):
         losses = []
         metrics = []
 
-        get_desc_str = (
-            lambda: f'Loss={np.mean(losses):.2f} Accuracy={np.mean(metrics):.2f}'
-        )
+        log_generator = LogGenerator('T', losses, metrics)
 
-        with ProgressBar(description, message_on_done=get_desc_str) as p:
+        with ProgressBar(
+            description, message_on_done=log_generator, final_line_feed=False
+        ) as p:
             for inputs, label in data:
                 # forward step
                 outputs = self.wrapped_model(*inputs)
@@ -91,7 +88,7 @@ def _train(self, data, optimizer: Optimizer, description: str):
                 losses.append(loss.item())
                 metrics.append(metric.numpy())
 
-                p.update(message=get_desc_str())
+                p.update(message=log_generator())
         return losses, metrics
 
     def fit(
@@ -128,7 +125,7 @@ def fit(
                     inputs=eval_data, batch_size=batch_size, shuffle=False
                 )
 
-                le, me = self._eval(_data)
+                le, me = self._eval(_data, train_log=LogGenerator('T', lt, mt)())
                 losses_eval.extend(le)
                 metrics_eval.extend(me)