Module Slicing

docs/api/Module.md

@@ -12,4 +12,5 @@
             - set_parameters
             - reset
             - init
+            - slice
 

docs/api/nn/EMAParamsTree.md

@@ -12,4 +12,5 @@
             - set_parameters
             - reset
             - init
+            - slice
 

docs/api/nn/Sequential.md

@@ -12,4 +12,5 @@
             - set_parameters
             - reset
             - init
+            - slice
 

elegy/__init__.py

@@ -30,6 +30,7 @@
     OutputStates,
     Mode,
 )
+from . import module_slicing
 from .generalized_module.generalized_module import GeneralizedModule
 from .generalized_optimizer.generalized_optimizer import GeneralizedOptimizer
 

elegy/hooks.py

@@ -105,6 +105,7 @@ def add_summary(
     path: types.Path,
     module: tp.Any,
     value: tp.Any,
+    input_values: tp.Optional[tp.Tuple[tp.Tuple, tp.Dict]] = None,
 ) -> None:
     """
     A hook that lets you define a summary in the current module. Its primary
@@ -127,8 +128,7 @@ def call(self, x):
 
     if not summaries_active():
         return
-
-    LOCAL.summaries.append(types.Summary(path, module, value))
+    LOCAL.summaries.append(types.Summary(path, module, value, input_values))
 
 
 def get_losses() -> types.Logs:

elegy/hooks_test.py

@@ -32,7 +32,7 @@ def test_summaries(self):
             elegy.hooks.add_summary(("a", 0, "b"), None, 2.0)
             summaries = elegy.hooks.get_summaries()
 
-        assert summaries[0] == (("a", 0, "b"), None, 2.0)
+        assert summaries[0] == (("a", 0, "b"), None, 2.0, None)
 
     def test_no_summaries(self):
         assert not elegy.hooks.summaries_active()
@@ -65,4 +65,4 @@ def f(x):
         assert x == 6
         assert losses["x_loss"] == 6
         assert metrics["x"] == 7
-        assert summaries[0] == (("a", 0, "b"), jax.nn.relu, 8)
+        assert summaries[0] == (("a", 0, "b"), jax.nn.relu, 8, None)

elegy/model/model.py

@@ -224,7 +224,7 @@ def summary_step(
 
         entries: tp.List[types.SummaryTableEntry] = []
 
-        for path, module, value in summaries:
+        for path, module, value, input_values in summaries:
 
             module_params, module_states = self.api_module.get_summary_params(
                 path=path,
@@ -240,7 +240,9 @@ def summary_step(
                     module_type_name=(
                         module.__class__.__name__ if is_generalizable(module) else ""
                     ),
+                    module=module,
                     output_value=value,
+                    input_value=input_values,
                     trainable_params_count=(
                         utils.parameters_count(module_params)
                         if module_params is not None

elegy/model/model_core.py

@@ -492,6 +492,7 @@ def summary(
         depth: int = 2,
         tablefmt: str = "fancy_grid",
         return_repr: bool = False,
+        return_raw_entries: bool = False,
         **tablulate_kwargs,
     ) -> tp.Optional[str]:
         """
@@ -529,6 +530,9 @@ def summary(
         total_entry = entries[-1]
         entries = entries[:-1]
 
+        if return_raw_entries:
+            return entries
+
         depth_groups: tp.Dict[str, tp.List[types.SummaryTableEntry]] = toolz.groupby(
             lambda entry: "/".join(entry.path.split("/")[:depth]), entries
         )
@@ -541,7 +545,9 @@ def get_grouped_entry(
             return types.SummaryTableEntry(
                 path=entry.path,
                 module_type_name=entry.module_type_name,
+                module=entry.module,
                 output_value=entry.output_value,
+                input_value=entry.input_value,
                 trainable_params_count=sum(
                     entry_.trainable_params_count for entry_ in group
                 ),

elegy/module.py

@@ -13,6 +13,11 @@
 from elegy import hooks, utils
 from elegy import types
 
+# placeholder for module module_slicing.py
+# injected from inside the module because of a circular dependency
+module_slicing = None
+
+
 __all__ = [
     "Module",
     "to_module",
@@ -158,6 +163,7 @@ class Module(metaclass=ModuleMeta):
         "set_parameters",
         "reset",
         "init",
+        "slice",
     ]
 
     def __init__(self, name: tp.Optional[str] = None, dtype: tp.Any = jnp.float32):
@@ -368,19 +374,25 @@ def __call__(self, *args, **kwargs) -> tp.Any:
             if hooks.summaries_active():
                 path = get_module_path(self)
                 assert path is not None
-                hooks.add_summary(path, self, outputs)
+                hooks.add_summary(path, self, outputs, (args, kwargs))
 
             return outputs
 
     @abstractmethod
     def call(self, *args, **kwargs):
         ...
 
-    def add_summary(self, name: str, f: tp.Any, value: tp.Any):
+    def add_summary(
+        self,
+        name: str,
+        f: tp.Any,
+        value: tp.Any,
+        input_values: tp.Optional[tp.Tuple[tp.Tuple, tp.Dict]] = None,
+    ):
         if hooks.summaries_active():
             path = get_module_path(self) + (name,)
             assert path is not None
-            hooks.add_summary(path, f, value)
+            hooks.add_summary(path, f, value, input_values)
 
     def init(
         self,
@@ -569,9 +581,54 @@ def add_parameter(
                 name=utils.get_name(regularizer),
                 value=regularizer(value),
             )
-
         return value
 
+    def slice(
+        self,
+        start_module: tp.Union["Module", str, None],
+        end_module: tp.Union[
+            "Module", str, None, tp.List[tp.Union["Module", str, None]]
+        ],
+        sample_input: np.ndarray,
+    ) -> "Module":
+        """
+        Creates a new submodule starting from the input of `start_module` to the outputs of `end_module`.
+
+        Current limitations:
+
+        - all operations between `start_module` and `end_module` must be performed by modules
+            i.e. `jax.nn.relu()` or `x+1` is not allowed but can be converted by wrapping with `elegy.to_module()`
+        - only one `start_module` is supported
+        - all modules between `start_module` and `end_module` must have a single output
+
+
+        Example usage:
+        ```
+        x = jnp.zeros((2, 224, 224, 3))
+        resnet = elegy.nets.resnet.ResNet18()
+        submodule = resnet.slice(
+                        start_module=None,
+                        end_module=["/res_net_block_1", "/res_net_block_3", "/res_net_block_5", "/res_net_block_7" ],
+                        sample_input=x,
+        )
+        outputs = elegy.Model(submodule).predict(x)
+        assert outputs[0].shape == (2, 56, 56, 64)
+        assert outputs[1].shape == (2, 28, 28, 128)
+        assert outputs[2].shape == (2, 14, 14, 256)
+        assert outputs[3].shape == (2, 7, 7, 512)
+        ```
+
+        Arguments:
+            start_module: Child module or name of a child module which will be the input module of the resulting module.
+                          If `None`, the first module is used.
+            end_module: Child module, name of child module, `None` or a list thereof which will be the output module(s) of the resulting module.
+                         If `None`, the last module is used.
+            sample_input: An array representing a sample input to the parent module.
+        """
+        return module_slicing.slice_module_from_to(
+            self, start_module, end_module, sample_input
+        )
+
     def update_parameter(self, name: str, value: tp.Any) -> None:
         """
         Update a parameter of the current module.