Adding Link

Author: marcromeyn

merlin/models/torch/block.py

@@ -22,6 +22,7 @@
 
 from merlin.models.torch.batch import Batch
 from merlin.models.torch.container import BlockContainer, BlockContainerDict
+from merlin.models.torch.link import Link, LinkType
 from merlin.models.torch.registry import registry
 from merlin.models.utils.registry import RegistryMixin
 
@@ -65,7 +66,7 @@ def forward(
 
         return inputs
 
-    def repeat(self, n: int = 1, name=None) -> "Block":
+    def repeat(self, n: int = 1, link: Optional[LinkType] = None, name=None) -> "Block":
         """
         Creates a new block by repeating the current block `n` times.
         Each repetition is a deep copy of the current block.
@@ -89,6 +90,9 @@ def repeat(self, n: int = 1, name=None) -> "Block":
             raise ValueError("n must be greater than 0")
 
         repeats = [self.copy() for _ in range(n - 1)]
+        if link:
+            parsed_link = Link.parse(link)
+            repeats = [parsed_link.copy().setup_link(repeat) for repeat in repeats]
 
         return Block(self, *repeats, name=name)
 
@@ -152,33 +156,33 @@ def forward(
 
         return outputs
 
-    def append(self, module: nn.Module):
-        self.post.append(module)
+    def append(self, module: nn.Module, link: Optional[LinkType] = None):
+        self.post.append(module, link=link)
 
         return self
 
-    def prepend(self, module: nn.Module):
-        self.pre.prepend(module)
+    def prepend(self, module: nn.Module, link: Optional[LinkType] = None):
+        self.pre.prepend(module, link=link)
 
         return self
 
-    def append_to(self, name: str, module: nn.Module):
-        self.branches[name].append(module)
+    def append_to(self, name: str, module: nn.Module, link: Optional[LinkType] = None):
+        self.branches[name].append(module, link=link)
 
         return self
 
-    def prepend_to(self, name: str, module: nn.Module):
-        self.branches[name].prepend(module)
+    def prepend_to(self, name: str, module: nn.Module, link: Optional[LinkType] = None):
+        self.branches[name].prepend(module, link=link)
 
         return self
 
-    def append_for_each(self, module: nn.Module, shared=False):
-        self.branches.append_for_each(module, shared=shared)
+    def append_for_each(self, module: nn.Module, shared=False, link: Optional[LinkType] = None):
+        self.branches.append_for_each(module, shared=shared, link=link)
 
         return self
 
-    def prepend_for_each(self, module: nn.Module, shared=False):
-        self.branches.prepend_for_each(module, shared=shared)
+    def prepend_for_each(self, module: nn.Module, shared=False, link: Optional[LinkType] = None):
+        self.branches.prepend_for_each(module, shared=shared, link=link)
 
         return self
 

merlin/models/torch/container.py

@@ -21,6 +21,7 @@
 from torch import nn
 from torch._jit_internal import _copy_to_script_wrapper
 
+from merlin.models.torch.link import Link, LinkType
 from merlin.models.torch.utils import torchscript_utils
 
 
@@ -46,7 +47,7 @@ def __init__(self, *inputs: nn.Module, name: Optional[str] = None):
 
         self._name: str = name
 
-    def append(self, module: nn.Module):
+    def append(self, module: nn.Module, link: Optional[Link] = None):
         """Appends a given module to the end of the list.
 
         Parameters
@@ -58,11 +59,12 @@ def append(self, module: nn.Module):
         -------
         self
         """
-        self.values.append(self.wrap_module(module))
+        _module = self._check_link(module, link=link)
+        self.values.append(self.wrap_module(_module))
 
         return self
 
-    def prepend(self, module: nn.Module):
+    def prepend(self, module: nn.Module, link: Optional[Link] = None):
         """Prepends a given module to the beginning of the list.
 
         Parameters
@@ -74,9 +76,9 @@ def prepend(self, module: nn.Module):
         -------
         self
         """
-        return self.insert(0, module)
+        return self.insert(0, module, link=link)
 
-    def insert(self, index: int, module: nn.Module):
+    def insert(self, index: int, module: nn.Module, link: Optional[Link] = None):
         """Inserts a given module at the specified index.
 
         Parameters
@@ -90,8 +92,8 @@ def insert(self, index: int, module: nn.Module):
         -------
         self
         """
-
-        self.values.insert(index, self.wrap_module(module))
+        _module = self._check_link(module, link=link)
+        self.values.insert(index, self.wrap_module(_module))
 
         return self
 
@@ -152,6 +154,15 @@ def __repr__(self) -> str:
     def _get_name(self) -> str:
         return super()._get_name() if self._name is None else self._name
 
+    def _check_link(self, module: nn.Module, link: Optional[LinkType] = None) -> nn.Module:
+        if link:
+            linked_module: Link = Link.parse(link)
+            linked_module.setup_link(module)
+
+            return linked_module
+
+        return module
+
 
 class BlockContainerDict(nn.ModuleDict):
     def __init__(
@@ -166,28 +177,36 @@ def __init__(
         super().__init__(modules)
         self._name: str = name
 
-    def append_to(self, name: str, module: nn.Module) -> "BlockContainerDict":
-        self._modules[name].append(module)
+    def append_to(
+        self, name: str, module: nn.Module, link: Optional[LinkType] = None
+    ) -> "BlockContainerDict":
+        self._modules[name].append(module, link=link)
 
         return self
 
-    def prepend_to(self, name: str, module: nn.Module) -> "BlockContainerDict":
-        self._modules[name].prepend(module)
+    def prepend_to(
+        self, name: str, module: nn.Module, link: Optional[LinkType] = None
+    ) -> "BlockContainerDict":
+        self._modules[name].prepend(module, link=link)
 
         return self
 
     # Append to all branches, optionally copying
-    def append_for_each(self, module: nn.Module, shared=False) -> "BlockContainerDict":
+    def append_for_each(
+        self, module: nn.Module, shared=False, link: Optional[LinkType] = None
+    ) -> "BlockContainerDict":
         for branch in self.values():
             _module = module if shared else deepcopy(module)
-            branch.append(_module)
+            branch.append(_module, link=link)
 
         return self
 
-    def prepend_for_each(self, module: nn.Module, shared=False) -> "BlockContainerDict":
+    def prepend_for_each(
+        self, module: nn.Module, shared=False, link: Optional[LinkType] = None
+    ) -> "BlockContainerDict":
         for branch in self.values():
             _module = module if shared else deepcopy(module)
-            branch.prepend(_module)
+            branch.prepend(_module, link=link)
 
         return self
 

merlin/models/torch/link.py

@@ -0,0 +1,83 @@
+import copy
+from typing import Dict, Optional, Union
+
+import torch
+from torch import nn
+
+from merlin.models.torch.registry import TorchRegistryMixin
+
+LinkType = Union[str, "Link"]
+
+
+class Link(nn.Module, TorchRegistryMixin):
+    """Base class for different types of network links.
+
+    This is typically used as part of a `Block` to connect different modules.
+
+    Some examples of links are:
+        - `residual`: Adds the input to the output of the module.
+        - `shortcut`: Outputs a dictionary with the output of the module and the input.
+        - `shortcut-concat`: Concatenates the input and the output of the module.
+
+    """
+
+    def __init__(self, output: Optional[nn.Module] = None):
+        super().__init__()
+
+        if output is not None:
+            self.setup_link(output)
+
+    def setup_link(self, output: nn.Module) -> "Link":
+        """
+        Setup function for the link.
+
+        Parameters
+        ----------
+        output : nn.Module
+            The output module for the link.
+
+        Returns
+        -------
+        Link
+            The updated Link instance.
+        """
+
+        self.output = output
+
+        return self
+
+    def copy(self) -> "Link":
+        """
+        Returns a copy of the link.
+
+        Returns
+        -------
+        Link
+            The copied link.
+        """
+        return copy.deepcopy(self)
+
+
+@Link.registry.register("residual")
+class Residual(Link):
+    """Adds the input to the output of the module."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.output(x)
+
+
+@Link.registry.register("shortcut")
+class Shortcut(Link):
+    """Outputs a dictionary with the output of the module and the input."""
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        return {"output": self.output(x), "shortcut": x}
+
+
+@Link.registry.register("shortcut-concat")
+class ShortcutConcat(Link):
+    """Concatenates the input and the output of the module."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        intermediate_output = self.output(x)
+        return torch.cat((x, intermediate_output), dim=1)

merlin/models/torch/registry.py

@@ -1,3 +1,10 @@
-from merlin.models.utils.registry import Registry
+from merlin.models.utils.registry import Registry, RegistryMixin
 
 registry: Registry = Registry.class_registry("modules")
+
+
+class TorchRegistryMixin(RegistryMixin):
+    registry = registry
+
+
+__all__ = ["registry", "Registry", "RegistryMixin", "TorchRegistryMixin"]

tests/unit/torch/test_block.py

@@ -19,6 +19,7 @@
 import torch
 from torch import nn
 
+from merlin.models.torch import link
 from merlin.models.torch.batch import Batch
 from merlin.models.torch.block import Block, ParallelBlock
 from merlin.models.torch.container import BlockContainer, BlockContainerDict
@@ -59,6 +60,9 @@ def test_insertion(self):
 
         assert torch.equal(outputs, inputs + 2)
 
+        block.append(PlusOne(), link="residual")
+        assert isinstance(block[-1], link.Residual)
+
     def test_copy(self):
         block = Block(PlusOne())
 
@@ -82,6 +86,19 @@ def test_repeat(self):
         with pytest.raises(ValueError, match="n must be greater than 0"):
             block.repeat(0)
 
+    def test_repeat_with_link(self):
+        block = Block(PlusOne())
+
+        repeated = block.repeat(2, link="residual")
+        assert isinstance(repeated, Block)
+        assert len(repeated) == 2
+        assert isinstance(repeated[-1], link.Residual)
+
+        inputs = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        outputs = module_utils.module_test(repeated, inputs)
+
+        assert torch.equal(outputs, (inputs + 1) + (inputs + 1) + 1)
+
     def test_from_registry(self):
         @Block.registry.register("my_block")
         class MyBlock(Block):