feat(tailor): add display and refactor summary (#112)

* feat(tailor): add display and refactor summary

* feat(tailor): add display and refactor summary

finetuner/helper.py

@@ -20,7 +20,7 @@
     Callable[..., DocumentSequence],
 ]  #: The type :py:data:`DocumentSequence` or a function that gives :py:data:`DocumentSequence`
 
-EmbeddingLayerInfoType = List[
+LayerInfoType = List[
     Dict[str, Any]
 ]  #: The type of embedding layer information used in Tailor
 TunerReturnType = Dict[

finetuner/tailor/__init__.py

@@ -30,3 +30,26 @@ def to_embedding_model(
     return ft(model, input_size, input_dtype).to_embedding_model(
         layer_name=layer_name, output_dim=output_dim, freeze=freeze
     )
+
+
+def display(
+    model: AnyDNN,
+    input_size: Optional[Tuple[int, ...]] = None,
+    input_dtype: str = 'float32',
+) -> AnyDNN:
+    f_type = get_framework(model)
+
+    if f_type == 'keras':
+        from .keras import KerasTailor
+
+        ft = KerasTailor
+    elif f_type == 'torch':
+        from .pytorch import PytorchTailor
+
+        ft = PytorchTailor
+    elif f_type == 'paddle':
+        from .paddle import PaddleTailor
+
+        ft = PaddleTailor
+
+    return ft(model, input_size, input_dtype).display()

finetuner/tailor/base.py

@@ -4,7 +4,7 @@
     Tuple,
 )
 
-from ..helper import AnyDNN, EmbeddingLayerInfoType
+from ..helper import AnyDNN, LayerInfoType
 
 
 class BaseTailor(abc.ABC):
@@ -49,10 +49,39 @@ def to_embedding_model(
         ...
 
     @property
-    @abc.abstractmethod
-    def embedding_layers(self) -> EmbeddingLayerInfoType:
+    def embedding_layers(self) -> LayerInfoType:
         """Get all dense layers that can be used as embedding layer from the :py:attr:`.model`.
 
-        :return: layers info as :class:`list` of :class:`dict`.
+        :return: layers info as Dict.
+        """
+        _layers = self.summary()
+        return [_l for _l in _layers if _l['is_embedding_layer']]
+
+    @abc.abstractmethod
+    def summary(self) -> LayerInfoType:
+        """The summary of the model architecture. To list all possible embedding layers, use :py:attr:`.embedding_layers`.
+
+        :return: layers info as Dict.
         """
         ...
+
+    def display(self) -> None:
+        """Display the model architecture from :py:attr:`.summary` in a table. """
+        from rich.table import Table
+        from rich import print, box
+
+        _summary = self.summary()
+        table = Table(box=box.SIMPLE)
+        cols = ['name', 'output_shape_display', 'nb_params', 'trainable']
+        for k in cols:
+            table.add_column(k)
+        for s in _summary:
+            style = None
+            if s['is_embedding_layer']:
+                style = 'green'
+            table.add_row(*map(str, (s[v] for v in cols)), style=style)
+        print(
+            table,
+            '[green]Green[/green] layers can be used as embedding layers, '
+            'whose [b]name[/b] can be used as [b]layer_name[/b] in to_embedding_model(...).',
+        )

finetuner/tailor/keras/__init__.py

@@ -1,23 +1,16 @@
 from typing import Optional
 
-from jina.helper import cached_property
 from tensorflow.keras import Model
 from tensorflow.keras.layers import Dense
 
 from ..base import BaseTailor
-from ...helper import EmbeddingLayerInfoType, AnyDNN
+from ...helper import LayerInfoType, AnyDNN
 
 
 class KerasTailor(BaseTailor):
     """Tailor class for Keras DNN models."""
 
-    @cached_property
-    def embedding_layers(self) -> EmbeddingLayerInfoType:
-        """Get all dense layers that can be used as embedding layer from the :py:attr:`.model`.
-
-        :return: layers info as :class:`list` of :class:`dict`.
-        """
-
+    def summary(self) -> LayerInfoType:
         def _get_shape(layer):
             try:
                 return layer.output_shape
@@ -27,31 +20,35 @@ def _get_shape(layer):
         results = []
         for idx, layer in enumerate(self._model.layers):
             output_shape = _get_shape(layer)
-            if (
+            is_embedding_layer = not (
                 not output_shape
                 or len(output_shape) != 2
                 or not isinstance(output_shape[-1], int)
-            ):
-                continue
-            else:
-                if not layer.built and not getattr(layer, '_is_graph_network', False):
-                    # If a subclassed model has a layer that is not called in Model.call, the
-                    # layer will not be built and we cannot call layer.count_params().
-                    params = 0
-                else:
-                    params = layer.count_params()
+            )
 
-                results.append(
-                    {
-                        'name': layer.name,
-                        'cls_name': layer.__class__.__name__,
-                        'output_shape': output_shape,
-                        'output_features': output_shape[-1],
-                        'nb_params': params,
-                        'layer_idx': idx,
-                        'module_name': layer.name,  # duplicate as `name` to make different backends consistent
-                    }
-                )
+            if not layer.built and not getattr(layer, '_is_graph_network', False):
+                # If a subclassed model has a layer that is not called in Model.call, the
+                # layer will not be built and we cannot call layer.count_params().
+                params = 0
+            else:
+                params = layer.count_params()
+
+            results.append(
+                {
+                    'name': layer.name,
+                    'cls_name': layer.__class__.__name__,
+                    'output_shape': output_shape,
+                    'output_shape_display': list(output_shape[1:]),
+                    'output_features': output_shape[
+                        -1
+                    ],  #: this only makes sense when is_embedding_layer is True
+                    'nb_params': params,
+                    'layer_idx': idx,
+                    'module_name': layer.name,  # duplicate as `name` to make different backends consistent
+                    'is_embedding_layer': is_embedding_layer,
+                    'trainable': layer.trainable,
+                }
+            )
         return results
 
     def to_embedding_model(

finetuner/tailor/paddle/__init__.py

@@ -1,15 +1,14 @@
 import copy
 import warnings
 from collections import OrderedDict
-from typing import Tuple, Optional
+from typing import Optional
 
 import numpy as np
 import paddle
-from jina.helper import cached_property
 from paddle import nn, Tensor
 
 from ..base import BaseTailor
-from ...helper import is_seq_int, EmbeddingLayerInfoType, AnyDNN
+from ...helper import is_seq_int, LayerInfoType, AnyDNN
 
 
 class PaddleTailor(BaseTailor):
@@ -19,12 +18,7 @@ class PaddleTailor(BaseTailor):
         To use this class, you need to set ``input_size`` and ``input_dtype`` in :py:meth:`.__init__`
     """
 
-    @cached_property
-    def embedding_layers(self) -> EmbeddingLayerInfoType:
-        """Get all dense layers that can be used as embedding layer from the :py:attr:`.model`.
-
-        :return: layers info as :class:`list` of :class:`dict`.
-        """
+    def summary(self) -> LayerInfoType:
         if not self._input_size:
             raise ValueError(
                 f'{self.__class__} requires a valid `input_size`, but receiving {self._input_size}'
@@ -69,6 +63,9 @@ def hook(layer, input, output):
                     params += np.prod(v.shape)
 
                 summary[m_key]['nb_params'] = params
+                summary[m_key]['trainable'] = any(
+                    l.trainable for _, l in layer_state_dict.items()
+                )
 
             if (
                 not isinstance(layer, nn.Sequential)
@@ -102,22 +99,22 @@ def hook(layer, input, output):
         results = []
         for idx, layer in enumerate(summary):
             output_shape = summary[layer]['output_shape']
-            if (
+            is_embedding_layer = not (
                 not output_shape
                 or len(output_shape) != 2
                 or not is_seq_int(output_shape)
                 or summary[layer]['cls_name'] in self._model.__class__.__name__
-            ):
-                continue
+            )
 
             results.append(
                 {
                     **summary[layer],
                     'output_features': output_shape[-1],
+                    'output_shape_display': output_shape[1:],
                     'layer_idx': idx,
+                    'is_embedding_layer': is_embedding_layer,
                 }
             )
-
         return results
 
     def to_embedding_model(

finetuner/tailor/pytorch/__init__.py

@@ -9,18 +9,13 @@
 from torch import nn
 
 from ..base import BaseTailor
-from ...helper import is_seq_int, EmbeddingLayerInfoType, AnyDNN
+from ...helper import is_seq_int, LayerInfoType, AnyDNN
 
 
 class PytorchTailor(BaseTailor):
     """Tailor class for PyTorch DNN models"""
 
-    @property
-    def embedding_layers(self) -> EmbeddingLayerInfoType:
-        """Get all dense layers that can be used as embedding layer from the :py:attr:`.model`.
-
-        :return: layers info as :class:`list` of :class:`dict`.
-        """
+    def summary(self) -> LayerInfoType:
         if not self._input_size:
             raise ValueError(
                 f'{self.__class__} requires a valid `input_size`, but receiving {self._input_size}'
@@ -54,6 +49,7 @@ def hook(module, input, output):
                 summary[m_key]['module_name'] = module.name
 
                 params = 0
+                summary[m_key]['trainable'] = False
                 if hasattr(module, 'weight') and hasattr(module.weight, 'size'):
                     params += np.prod(list(module.weight.size()))
                     summary[m_key]['trainable'] = module.weight.requires_grad
@@ -89,19 +85,20 @@ def hook(module, input, output):
         results = []
         for idx, layer in enumerate(summary):
             output_shape = summary[layer]['output_shape']
-            if (
+            is_embedding_layer = not (
                 not output_shape
                 or len(output_shape) != 2
                 or not is_seq_int(output_shape)
                 or summary[layer]['cls_name'] in self._model.__class__.__name__
-            ):
-                continue
+            )
 
             results.append(
                 {
                     **summary[layer],
                     'output_features': output_shape[-1],
+                    'output_shape_display': output_shape[1:],
                     'layer_idx': idx,
+                    'is_embedding_layer': is_embedding_layer,
                 }
             )
 

tests/unit/tailor/test_common.py

@@ -4,7 +4,7 @@
 import torch
 
 from finetuner.helper import get_framework
-from finetuner.tailor import to_embedding_model
+from finetuner.tailor import to_embedding_model, display
 
 
 class LastCellPT(torch.nn.Module):
@@ -53,3 +53,9 @@ def test_to_embedding_fn(framework, output_dim, freeze):
     )
     assert m1
     assert get_framework(m1) == framework
+
+
+@pytest.mark.parametrize('framework', ['keras', 'paddle', 'torch'])
+def test_display(framework):
+    m = embed_models[framework]()
+    display(m, input_size=(5000,), input_dtype='int64')