Repository: pystiche/pystiche

Issue:
Improve multi-layer encoding
Currently the canonical way for multi-layer encoding is to call `enc.MultiLayerEncoder.encode()` [**before** the criterion is called](https://github.com/pmeier/pystiche/blob/ad9b3ea09f9803249b162fbe009db25869e42fb4/pystiche/loss/multi_op.py#L61-L66). This ensures that every encoding is present when requested by an `ops.EncodingOperator` in the criterion.

Given that the layers are registered anyway this extra step is not needed: we could improve `enc.MultiLayerEncoder.forward()` to resume encoding from a previously stored encoding from the same pass.

Example:

```python
from pystiche import enc, ops

multi_layer_encoder = enc.vgg19_multi_layer_encoder(pretrained=False)

op_shallow = ops.FeatureReconstructionOperator(
    multi_layer_encoder.extract_encoder("relu2_1")
)
op_deep = ops.FeatureReconstructionOperator(
    multi_layer_encoder.extract_encoder("relu5_1")
)
```

Calling it from shallow to deep

```python
op_shallow(input)
op_deep(input)
```

should now encode up to `relu2_1` in the first call and pick up from there up to `relu5_1` in the second call. The other way around

```python
op_deep(input)
op_shallow(input)
```

the first call should encode up to `relu5_1` storing the encodings of  `relu2_1` in the process. Afterwards the second call is just returning the encodings from storage.

Thus, although it is not a single forward pass any more, it has the same cost as before but removes the need to call `enc.MultiLayerEncoder.encode()`. For the (unlikely) case that the user does not need the encodings of all layers in every step this might also reduce the computation costs.


Discussion:
Upon further inspection, I think this will be a BC breaking change. Right now you call an `enc.MultiLayerEncoder` to get the encodings of the input and invoke `encode()` to store them for later use. I think it would be better swap this functionality: 

With the addition detailed in the first comment, calling `enc.MultiLayerEncoder` should happen with a single layer and the result should be stored by default. Invoking `encode` on the other hand should simply return the encodings for the given input.

Resolving patch:
diff --git a/pystiche/enc/multi_layer_encoder.py b/pystiche/enc/multi_layer_encoder.py
--- a/pystiche/enc/multi_layer_encoder.py
+++ b/pystiche/enc/multi_layer_encoder.py
@@ -1,173 +1,253 @@
 import warnings
-from collections import OrderedDict
-from copy import copy
-from typing import Collection, Dict, Iterator, Optional, Sequence, Set, Tuple, cast
+from collections import OrderedDict, defaultdict
+from typing import (
+    Any,
+    Callable,
+    Collection,
+    DefaultDict,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    cast,
+)
 
 import torch
 from torch import nn
 
 import pystiche
-from pystiche.misc import suppress_warnings
 
+from ..misc import build_deprecation_message
 from .encoder import Encoder
 from .guides import propagate_guide
 
 __all__ = ["MultiLayerEncoder", "SingleLayerEncoder"]
 
 
-def _future_warning(name: str) -> None:
-    msg = (
-        f"The functionality of MultiLayerEncoder.{name} will change in the future. "
-        f"If you depend on this functionality, "
-        f"see https://github.com/pmeier/pystiche/issues/435 for details "
-    )
-    warnings.warn(msg, FutureWarning)
+class _Layers:
+    def __init__(self, modules: Dict[str, nn.Module]) -> None:
+        self._modules = modules
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._modules
+
+    def __len__(self) -> int:
+        return len(self._modules)
+
+    @property
+    def _names(self) -> Tuple[str, ...]:
+        # TODO: Check if tuple generation is expensive. If that is the case, cache it
+        #  based on self._modules.keys()
+        return tuple(self._modules.keys())
+
+    def _name_to_idx(self, name: str) -> int:
+        if name not in self:
+            raise ValueError
+
+        return self._names.index(name)
+
+    def _idx_to_name(self, idx: int) -> str:
+        if not (0 <= idx < len(self)):
+            raise ValueError
+
+        return self._names[idx]
+
+    def range(
+        self,
+        start: Optional[str] = None,
+        stop: Optional[str] = None,
+        include_start: bool = True,
+        include_stop: bool = True,
+    ) -> Tuple[str, ...]:
+        if not (start or stop):
+            return self._names
+
+        if start is None:
+            start_idx = 0
+        else:
+            start_idx = self._name_to_idx(start)
+            if not include_start:
+                start_idx += 1
+
+        if stop is None:
+            stop_idx = len(self)
+        else:
+            stop_idx = self._name_to_idx(stop)
+            if include_stop:
+                stop_idx += 1
+
+        return self._names[start_idx:stop_idx]
+
+    def _depth(
+        self, names: Iterable[str], extractor: Callable[[List[int]], int]
+    ) -> str:
+        return self._idx_to_name(extractor([self._name_to_idx(name) for name in names]))
+
+    def shallowest(self, names: Optional[Iterable[str]] = None) -> str:
+        return self._depth(names or self._names, min)
+
+    def deepest(self, names: Optional[Iterable[str]] = None) -> str:
+        return self._depth(names or self._names, max)
+
+    def _neighbour(
+        self,
+        name: str,
+        names: Iterable[str],
+        edge_idx: int,
+        extractor: Callable[[int, List[int]], Optional[str]],
+    ) -> Optional[str]:
+        if not names:
+            return None
+
+        idx = self._name_to_idx(name)
+        idcs = [self._name_to_idx(name) for name in names]
+
+        if edge_idx in idcs:
+            return self._idx_to_name(edge_idx)
+
+        return extractor(idx, idcs)
+
+    def _extract_prev(self, idx: int, idcs: List[int]) -> Optional[str]:
+        candidates = [other_idx for other_idx in idcs if other_idx < idx]
+        if not candidates:
+            return None
+
+        return self._idx_to_name(max(candidates))
+
+    def prev(self, name: str, names: Iterable[str]) -> Optional[str]:
+        return self._neighbour(name, names, edge_idx=0, extractor=self._extract_prev)
+
+    def _extract_next(self, idx: int, idcs: List[int]) -> Optional[str]:
+        candidates = [other_idx for other_idx in idcs if other_idx > idx]
+        if not candidates:
+            return None
+
+        return self._idx_to_name(min(candidates))
+
+    def next(self, name: str, names: Collection[str]) -> Optional[str]:
+        return self._neighbour(
+            name, names, edge_idx=len(self) - 1, extractor=self._extract_next
+        )
 
 
 class MultiLayerEncoder(pystiche.Module):
-    r"""Sequential architecture that supports extracting encodings from intermediate
-    layers in a single forward pass. Encodings can be stored to avoid recalculating
-    them if they are required multiple times. Invokes :meth:`MultiLayerEncoder.forward`
-    if called.
+    r"""Sequential encoder with convenient access to intermediate layers.
 
     Args:
-        modules: Named sequential modules.
+        modules: Named modules that serve as basis for the encoding.
 
     Attributes:
-        registered_layers: Set of registered layers for
-            :meth:`MultiLayerEncoder.encode` and :meth:`MultiLayerEncoder.trim`.
+        registered_layers: Layers, on which the encodings will be cached during the
+            :meth:`forward` pass.
     """
 
     def __init__(self, modules: Sequence[Tuple[str, nn.Module]]) -> None:
         super().__init__(named_children=modules)
+        self._layers: _Layers = _Layers(self._modules)
         self.registered_layers: Set[str] = set()
-        self._storage: Dict[Tuple[str, pystiche.TensorKey], torch.Tensor] = dict()
-
-        # TODO: remove this?
-        self.requires_grad_(False)
-        self.eval()
-
-    def children_names(self) -> Iterator[str]:
-        for name, child in self.named_children():
-            yield name
+        self._cache: DefaultDict[
+            pystiche.TensorKey, Dict[str, torch.Tensor]
+        ] = defaultdict(lambda: {})
 
     def __contains__(self, layer: str) -> bool:
-        r"""Checks if the given layer is part of the :class:`MultiLayerEncoder`
+        r"""Is the layer part of the multi-layer encoder?
 
         Args:
-            layer: Layer.
+            layer: Layer to be checked.
         """
-        return layer in self.children_names()
-
-    def _verify_layer(self, layer: str) -> None:
-        if layer not in self:
-            raise ValueError(f"Layer {layer} is not part of the encoder.")
-
-    def extract_deepest_layer(self, layers: Collection[str]) -> str:
-        for layer in layers:
-            self._verify_layer(layer)
-        return sorted(set(layers), key=list(self.children_names()).index)[-1]
-
-    def named_children_to(
-        self, layer: str, include_last: bool = False
-    ) -> Iterator[Tuple[str, nn.Module]]:
-        self._verify_layer(layer)
-        idx = list(self.children_names()).index(layer)
-        if include_last:
-            idx += 1
-        for name, child in tuple(self.named_children())[:idx]:
-            yield name, child
-
-    def named_children_from(
-        self, layer: str, include_first: bool = True
-    ) -> Iterator[Tuple[str, nn.Module]]:
-        self._verify_layer(layer)
-        idx = list(self.children_names()).index(layer)
-        if not include_first:
-            idx += 1
-        for name, child in tuple(self.named_children())[idx:]:
-            yield name, child
+        return layer in self._layers
 
-    def forward(
-        self, input: torch.Tensor, layers: Sequence[str], store: bool = False
-    ) -> Tuple[torch.Tensor, ...]:
-        r"""Encode the input on the given layers in a single forward pass. If the input
-        was encoded before the encodings are extracted from the storage rather than
-        executing the forward pass again.
+    def _verify(self, name: str) -> None:
+        if name not in self:
+            raise ValueError(f"Layer {name} is not part of the multi-layer encoder.")
 
-        Args:
-            input: Input.
-            layers: Layers.
-            store: If ``True``, store the encodings.
+    def register_layer(self, layer: str) -> None:
+        r"""Register a layer for caching the encodings in the :meth:`forward` pass.
 
-        Returns:
-            Tuple of encodings which order corresponds to ``layers``.
+        Args:
+            layer: Layer to be registered.
         """
-        _future_warning("__call__")
-        storage = copy(self._storage)
-        input_key = pystiche.TensorKey(input)
-        stored_layers = [name for name, key in storage.keys() if key == input_key]
-        diff_layers = set(layers) - set(stored_layers)
-
-        if diff_layers:
-            deepest_layer = self.extract_deepest_layer(diff_layers)
-            for name, module in self.named_children_to(
-                deepest_layer, include_last=True
-            ):
-                input = storage[(name, input_key)] = module(input)
-
-            if store:
-                self._storage = storage
+        self._verify(layer)
+        self.registered_layers.add(layer)
 
-        return tuple(storage[(name, input_key)] for name in layers)
+    # FIXME: could this be moved into pystiche.Module?
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        r"""Invokes :meth:`forward`."""
+        return super().__call__(*args, **kwargs)
 
-    def extract_encoder(self, layer: str) -> "SingleLayerEncoder":
-        r"""Extract a :class:`SingleLayerEncoder` for the given layer and register
-        the layer in :attr:`MultiLayerEncoder.registered_layers`.
+    def forward(
+        self,
+        input: torch.Tensor,
+        layer: str,
+        cache: Optional[Dict[str, torch.Tensor]] = None,
+        to_cache: Optional[Collection[str]] = None,
+    ) -> torch.Tensor:
+        r"""Encode the input on layer.
 
         Args:
-            layer: Layer.
+            input: Input to be encoded.
+            layer: Layer on which the ``input`` should be encoded.
+            cache: Encoding cache. If omitted, defaults to the the internal cache.
+            to_cache: Layers, of which the encodings should be cached. If omitted,
+                defaults to :attr:`registered_layers`.
+
+        Examples:
+
+            >>> modules = [("conv", nn.Conv2d(3, 3, 3)), ("pool", nn.MaxPool2d(2))]
+            >>> mle = enc.MultiLayerEncoder(modules)
+            >>> input = torch.rand(1, 3, 128, 128)
+            >>> output = mle(input, "conv")
         """
-        self._verify_layer(layer)
-        self.registered_layers.add(layer)
-        return SingleLayerEncoder(self, layer)
+        self._verify(layer)
 
-    def encode(self, input: torch.Tensor) -> None:
-        r"""Encode the given input and store the encodings of all
-        :attr:`MultiLayerEncoder.registered_layers`.
+        if cache is None:
+            cache = self._cache[pystiche.TensorKey(input)]
+        if layer in cache:
+            return cache[layer]
 
-        Args:
-            input: Input.
-        """
-        _future_warning("encode")
+        if to_cache is None:
+            to_cache = self.registered_layers
+
+        prev = self._layers.prev(layer, cache.keys())
+        if prev is not None:
+            input = cache[prev]
+
+        for name in self._layers.range(prev, layer, include_start=False):
+            module = self._modules[name]
+            input = module(input)
+
+            if name in to_cache:
+                cache[name] = input
 
-        if not self.registered_layers:
-            return
+        return input
 
-        key = pystiche.TensorKey(input)
-        keys = [(layer, key) for layer in self.registered_layers]
-        encs = self(input, layers=self.registered_layers, store=True)
-        self._storage = dict(zip(keys, encs))
+    def clear_cache(self) -> None:
+        r"""Clear the internal cache."""
+        self._cache = defaultdict(lambda: {})
 
     def empty_storage(self) -> None:
-        r"""Empty the encodings storage."""
-        self._storage = {}
+        msg = build_deprecation_message(
+            "The method 'empty_storage'", "1.0", info="It was renamed to 'clear_cache'."
+        )
+        warnings.warn(msg)
+        self.clear_cache()
 
-    def trim(self, layers: Optional[Collection[str]] = None) -> None:
-        r"""Remove excess layers that are not necessary to generate the encodings of
-        ``layers``.
+    def encode(
+        self, input: torch.Tensor, layers: Sequence[str],
+    ) -> Tuple[torch.Tensor, ...]:
+        r"""Encode the input on layers.
 
         Args:
-            layers: Layers that the :class:`MultiLayerEncoder` need to be able to
-                generate encodings for. If ``None``,
-                :attr:`MultiLayerEncoder.registered_layers` is used. Defaults to
-                ``None``.
+            input: Input to be encoded.
+            layers: Layers on which the ``input`` should be encoded.
         """
-        if layers is None:
-            layers = self.registered_layers
-        deepest_layer = self.extract_deepest_layer(layers)
-        for name, _ in self.named_children_from(deepest_layer, include_first=False):
-            del self._modules[name]
+        cache: Dict[str, torch.Tensor] = {}
+        return tuple(
+            self(input, layer, cache=cache, to_cache=layers) for layer in layers
+        )
 
     def propagate_guide(
         self,
@@ -188,8 +268,8 @@ def propagate_guide(
             Tuple of guides which order corresponds to ``layers``.
         """
         guides = {}
-        deepest_layer = self.extract_deepest_layer(layers)
-        for name, module in self.named_children_to(deepest_layer, include_last=True):
+        for name in self._layers.range(stop=self._layers.deepest(layers)):
+            module = self._modules[name]
             try:
                 guide = guides[name] = propagate_guide(
                     module, guide, method=method, allow_empty=allow_empty
@@ -201,6 +281,27 @@ def propagate_guide(
 
         return tuple(guides[name] for name in layers)
 
+    def trim(self, layers: Optional[Iterable[str]] = None) -> None:
+        if layers is None:
+            layers = self.registered_layers
+        else:
+            for name in layers:
+                self._verify(name)
+
+        for name in self._layers.range(
+            self._layers.deepest(layers), include_start=False
+        ):
+            del self._modules[name]
+
+    def extract_encoder(self, layer: str) -> "SingleLayerEncoder":
+        r"""Extract a :class:`SingleLayerEncoder` for the layer and register it.
+
+        Args:
+            layer: Layer.
+        """
+        self.register_layer(layer)
+        return SingleLayerEncoder(self, layer)
+
 
 class SingleLayerEncoder(Encoder):
     r"""Encoder extracted from a :class:`MultiLayerEncoder` that operates on a single
@@ -216,18 +317,14 @@ def __init__(self, multi_layer_encoder: MultiLayerEncoder, layer: str):
         self.multi_layer_encoder = multi_layer_encoder
         self.layer = layer
 
-    def forward(self, input_image: torch.Tensor) -> torch.Tensor:
-        r"""Encode the given input image on :attr:`SingleLayerEncoder.layer` of
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        r"""Encode the given input on :attr:`SingleLayerEncoder.layer` of
         :attr:`SingleLayerEncoder.multi_layer_encoder`.
 
         Args:
             input_image: Input image.
         """
-        with suppress_warnings(FutureWarning):
-            return cast(
-                Tuple[torch.Tensor],
-                self.multi_layer_encoder(input_image, layers=(self.layer,)),
-            )[0]
+        return cast(torch.Tensor, self.multi_layer_encoder(input, self.layer))
 
     def propagate_guide(self, guide: torch.Tensor) -> torch.Tensor:
         r"""Propagate the given guide on :attr:`SingleLayerEncoder.layer` of
diff --git a/pystiche/loss/multi_op.py b/pystiche/loss/multi_op.py
--- a/pystiche/loss/multi_op.py
+++ b/pystiche/loss/multi_op.py
@@ -1,3 +1,4 @@
+import warnings
 from types import TracebackType
 from typing import Iterator, Sequence, Tuple, Type
 
@@ -6,7 +7,7 @@
 
 import pystiche
 from pystiche import enc, ops
-from pystiche.misc import suppress_warnings
+from pystiche.misc import build_deprecation_message
 
 __all__ = ["MLEHandler", "MultiOperatorLoss"]
 
@@ -21,21 +22,32 @@ def __init__(self, criterion: nn.Module) -> None:
         }
 
     def encode(self, input_image: torch.Tensor) -> None:
-        for encoder in self.multi_layer_encoders:
-            encoder.encode(input_image)
+        msg = build_deprecation_message(
+            "The method 'encode'",
+            "1.0",
+            info=(
+                "It is no longer needed to pre-encode the input. "
+                "See https://github.com/pmeier/pystiche/issues/435 for details"
+            ),
+        )
+        warnings.warn(msg)
+
+    def clear_cache(self) -> None:
+        for mle in self.multi_layer_encoders:
+            mle.clear_cache()
 
     def empty_storage(self) -> None:
-        for mle in self.multi_layer_encoders:
-            mle.empty_storage()
+        msg = build_deprecation_message(
+            "The method 'empty_storage'", "1.0", info="It was renamed to 'clear_cache'."
+        )
+        warnings.warn(msg)
+        self.clear_cache()
 
     def trim(self) -> None:
         for mle in self.multi_layer_encoders:
             mle.trim()
 
     def __call__(self, input_image: torch.Tensor) -> "MLEHandler":
-        with suppress_warnings(FutureWarning):
-            for mle in self.multi_layer_encoders:
-                mle.encode(input_image)
         return self
 
     def __enter__(self) -> None:
@@ -45,7 +57,7 @@ def __exit__(
         self, exc_type: Type[Exception], exc_val: Exception, exc_tb: TracebackType
     ) -> None:
         for encoder in self.multi_layer_encoders:
-            encoder.empty_storage()
+            encoder.clear_cache()
 
     def _named_children(self) -> Iterator[Tuple[str, enc.MultiLayerEncoder]]:
         return ((str(idx), mle) for idx, mle in enumerate(self.multi_layer_encoders))


Test patch:
diff --git a/tests/fixtures.py b/tests/fixtures.py
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -6,7 +6,7 @@
 
 import pystiche
 
-from . import assets, utils
+from . import assets, mocks, utils
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -63,3 +63,16 @@ def forward(self, input):
             return input
 
     return ForwardPassCounter()
+
+
+@pytest.fixture
+def module_factory():
+    def factory():
+        return mocks.ModuleMock()
+
+    return factory
+
+
+@pytest.fixture
+def module(module_factory):
+    return module_factory()
diff --git a/tests/integration/enc/test_multi_layer_encoder.py b/tests/integration/enc/test_multi_layer_encoder.py
--- a/tests/integration/enc/test_multi_layer_encoder.py
+++ b/tests/integration/enc/test_multi_layer_encoder.py
@@ -6,280 +6,216 @@
 
 from pystiche import enc
 
-from tests.asserts import assert_named_modules_identical
+from tests import mocks
 
 
-def test_MultiLayerEncoder():
-    modules = [(str(idx), nn.Module()) for idx in range(3)]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
-
-    for name, module in modules:
-        actual = getattr(multi_layer_encoder, name)
-        desired = module
-        assert actual is desired
+@pytest.fixture
+def mle_and_modules(module_factory):
+    shallow = module_factory()
+    intermediate = module_factory()
+    deep = module_factory()
+    modules = [("shallow", shallow), ("intermediate", intermediate), ("deep", deep)]
+    mle = enc.MultiLayerEncoder(modules)
+    return mle, dict(modules)
 
 
-def test_MultiLayerEncoder_named_children():
-    modules = [(str(idx), nn.Module()) for idx in range(3)]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+@pytest.fixture
+def mle(mle_and_modules):
+    return mle_and_modules[0]
 
-    actual = tuple(multi_layer_encoder.children_names())
-    desired = tuple(zip(*modules))[0]
-    assert actual == desired
 
+@pytest.fixture
+def input():
+    torch.manual_seed(0)
+    return torch.rand(1, 3, 16, 16)
 
-def test_MultiLayerEncoder_contains():
-    idcs = (0, 2)
-    modules = [(str(idx), nn.Module()) for idx in idcs]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
 
-    for idx in idcs:
-        assert str(idx) in multi_layer_encoder
+def test_MultiLayerEncoder_contains(mle):
+    assert "shallow" in mle
+    assert "intermediate" in mle
+    assert "deep" in mle
+    assert "unknown" not in mle
 
-    for idx in set(range(max(idcs))) - set(idcs):
-        assert str(idx) not in multi_layer_encoder
 
+def test_MultiLayerEncoder_register_layer(mle):
+    layer = "shallow"
 
-def test_MultiLayerEncoder_extract_deepest_layer():
-    layers = [str(idx) for idx in range(3)]
-    modules = [(layer, nn.Module()) for layer in layers]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+    assert layer not in mle.registered_layers
 
-    actual = multi_layer_encoder.extract_deepest_layer(layers)
-    desired = layers[-1]
-    assert actual == desired
+    mle.register_layer(layer)
 
-    actual = multi_layer_encoder.extract_deepest_layer(sorted(layers, reverse=True))
-    desired = layers[-1]
-    assert actual == desired
+    assert layer in mle.registered_layers
 
-    del multi_layer_encoder._modules[layers[-1]]
 
+def test_MultiLayerEncoder_register_layer_error(mle):
     with pytest.raises(ValueError):
-        multi_layer_encoder.extract_deepest_layer(layers)
+        mle.register_layer("unknown")
 
-    layers = layers[:-1]
 
-    actual = multi_layer_encoder.extract_deepest_layer(layers)
-    desired = layers[-1]
-    assert actual == desired
+def test_MultiLayerEncoder_call(input):
+    conv = nn.Conv2d(3, 1, 1)
+    relu = nn.ReLU(inplace=False)
 
+    modules = (("conv", conv), ("relu", relu))
+    mle = enc.MultiLayerEncoder(modules)
 
-def test_MultiLayerEncoder_named_children_to():
-    layers = [str(idx) for idx in range(3)]
-    modules = [(layer, nn.Module()) for layer in layers]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+    actual = mle(input, "conv")
+    expected = conv(input)
+    ptu.assert_allclose(actual, expected)
 
-    actuals = multi_layer_encoder.named_children_to(layers[-2])
-    desireds = modules[:-2]
-    assert_named_modules_identical(actuals, desireds)
+    actual = mle(input, "relu")
+    expected = relu(conv(input))
+    ptu.assert_allclose(actual, expected)
 
-    actuals = multi_layer_encoder.named_children_to(layers[-2], include_last=True)
-    desireds = modules[:-1]
-    assert_named_modules_identical(actuals, desireds)
 
+def test_MultiLayerEncoder_call_error(mle):
+    with pytest.raises(ValueError):
+        mle(torch.empty(1), "unknown")
 
-def test_MultiLayerEncoder_named_children_from():
-    layers = [str(idx) for idx in range(3)]
-    modules = [(layer, nn.Module()) for layer in layers]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
 
-    actuals = multi_layer_encoder.named_children_from(layers[-2])
-    desireds = modules[1:]
-    assert_named_modules_identical(actuals, desireds)
+def test_MultiLayerEncoder_call_use_cached(mle_and_modules, input):
+    mle, modules = mle_and_modules
+    layer = "deep"
+    mle.register_layer(layer)
 
-    actuals = multi_layer_encoder.named_children_from(layers[-2], include_first=False)
-    desireds = modules[2:]
-    assert_named_modules_identical(actuals, desireds)
+    mle(input, layer)
+    modules[layer].assert_called_once_with(input)
 
+    mle(input, layer)
+    modules[layer].assert_called_once()
 
-def test_MultiLayerEncoder_call():
-    torch.manual_seed(0)
-    conv = nn.Conv2d(3, 1, 1)
-    relu = nn.ReLU(inplace=False)
-    pool = nn.MaxPool2d(2)
-    input = torch.rand(1, 3, 128, 128)
 
-    modules = (("conv", conv), ("relu", relu), ("pool", pool))
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+def test_MultiLayerEncoder_call_resume_from_cache(mle_and_modules, input):
+    mle, modules = mle_and_modules
+    mle.register_layer("shallow")
 
-    layers = ("conv", "pool")
-    encs = multi_layer_encoder(input, layers)
+    mle(input, "shallow")
+    modules["shallow"].assert_called_once_with(input)
 
-    actual = encs[0]
-    desired = conv(input)
-    ptu.assert_allclose(actual, desired)
+    mle(input, "deep")
+    modules["shallow"].assert_called_once()
+    modules["deep"].assert_called_once_with(input)
 
-    actual = encs[1]
-    desired = pool(relu(conv(input)))
-    ptu.assert_allclose(actual, desired)
 
+def test_MultiLayerEncoder_call_pre_cache(mle_and_modules, input):
+    mle, modules = mle_and_modules
+    mle.register_layer("shallow")
 
-def test_MultiLayerEncoder_call_store(forward_pass_counter):
-    torch.manual_seed(0)
+    mle(input, "deep")
+    modules["shallow"].assert_called_once_with(input)
+    modules["deep"].assert_called_once_with(input)
 
-    modules = (("count", forward_pass_counter),)
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+    mle(input, "shallow")
+    modules["shallow"].assert_called_once()
 
-    layers = ("count",)
-    input = torch.rand(1, 3, 128, 128)
-    multi_layer_encoder(input, layers, store=True)
-    multi_layer_encoder(input, layers)
 
-    actual = forward_pass_counter.count
-    desired = 1
-    assert actual == desired
+def test_MultiLayerEncoder_clear_cache(mle_and_modules, input):
+    mle, modules = mle_and_modules
+    layer = "shallow"
+    module = modules[layer]
+    mle.register_layer(layer)
 
-    new_input = torch.rand(1, 3, 128, 128)
-    multi_layer_encoder(new_input, layers)
+    mle(input, layer)
+    module.assert_called_once_with(input)
 
-    actual = forward_pass_counter.count
-    desired = 2
-    assert actual == desired
+    mle.clear_cache()
 
+    mle(input, layer)
+    module.assert_called(2)
+    module.assert_called_with(input)
 
-def test_MultiLayerEncoder_extract_encoder():
-    conv = nn.Conv2d(3, 1, 1)
-    relu = nn.ReLU(inplace=False)
 
-    modules = (("conv", conv), ("relu", relu))
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+def test_MultiLayerEncoder_empty_storage(mocker, mle):
+    mock = mocker.patch(
+        mocks.make_mock_target(
+            "enc", "multi_layer_encoder", "MultiLayerEncoder", "clear_cache"
+        )
+    )
 
-    layer = "relu"
-    single_layer_encoder = multi_layer_encoder.extract_encoder(layer)
+    with pytest.warns(UserWarning):
+        mle.empty_storage()
 
-    assert isinstance(single_layer_encoder, enc.SingleLayerEncoder)
-    assert single_layer_encoder.multi_layer_encoder is multi_layer_encoder
-    assert single_layer_encoder.layer == layer
-    assert layer in multi_layer_encoder.registered_layers
+    mock.assert_called()
 
 
-def test_MultiLayerEncoder_encode(forward_pass_counter):
-    torch.manual_seed(0)
+def test_MultiLayerEncoder_encode(input):
     conv = nn.Conv2d(3, 1, 1)
     relu = nn.ReLU(inplace=False)
-    input = torch.rand(1, 3, 128, 128)
 
-    modules = (("count", forward_pass_counter), ("conv", conv), ("relu", relu))
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+    modules = (("conv", conv), ("relu", relu))
+    mle = enc.MultiLayerEncoder(modules)
 
-    layers = ("conv", "relu")
-    multi_layer_encoder.registered_layers.update(layers)
-    multi_layer_encoder.encode(input)
-    encs = multi_layer_encoder(input, layers)
+    encs = mle.encode(input, ("conv", "relu"))
 
     actual = encs[0]
-    desired = conv(input)
-    ptu.assert_allclose(actual, desired)
+    expected = conv(input)
+    ptu.assert_allclose(actual, expected)
 
     actual = encs[1]
-    desired = relu(conv(input))
-    ptu.assert_allclose(actual, desired)
-
-    actual = forward_pass_counter.count
-    desired = 1
-    assert actual == desired
-
-
-def test_MultiLayerEncoder_empty_storage(forward_pass_counter):
-    torch.manual_seed(0)
-    input = torch.rand(1, 3, 128, 128)
-
-    modules = (("count", forward_pass_counter),)
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+    expected = relu(conv(input))
+    ptu.assert_allclose(actual, expected)
 
-    layers = ("count",)
-    multi_layer_encoder(input, layers, store=True)
-    multi_layer_encoder.empty_storage()
-    multi_layer_encoder(input, layers)
 
-    actual = forward_pass_counter.count
-    desired = 2
-    assert actual == desired
+@pytest.mark.parametrize(
+    "layers",
+    (("shallow", "deep"), ("deep", "shallow")),
+    ids=lambda layers: f"{layers[0]} to {layers[1]}",
+)
+def test_MultiLayerEncoder_encode_cache(mle_and_modules, input, layers):
+    mle, modules = mle_and_modules
 
+    mle.encode(input, layers)
 
-def test_MultiLayerEncoder_trim():
-    layers = [str(idx) for idx in range(3)]
-    modules = [(layer, nn.Module()) for layer in layers]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
-
-    for name, module in modules:
-        actual = getattr(multi_layer_encoder, name)
-        desired = module
-        assert actual is desired
+    modules["shallow"].assert_called_once_with(input)
+    modules["deep"].assert_called_once_with(input)
 
-    idx = 1
-    multi_layer_encoder.trim((str(idx),))
 
-    for name, module in modules[: idx + 1]:
-        actual = getattr(multi_layer_encoder, name)
-        desired = module
-        assert actual is desired
+def test_MultiLayerEncoder_trim(mle):
+    assert "shallow" in mle
+    assert "intermediate" in mle
+    assert "deep" in mle
 
-    for name in tuple(zip(*modules))[0][idx + 1 :]:
-        with pytest.raises(AttributeError):
-            getattr(multi_layer_encoder, name)
+    mle.trim(("shallow",))
 
+    assert "shallow" in mle
+    assert "intermediate" not in mle
+    assert "deep" not in mle
 
-def test_MultiLayerEncoder_trim_layers():
-    layers = [str(idx) for idx in range(3)]
-    modules = [(layer, nn.Module()) for layer in layers]
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
 
-    for name, module in modules:
-        actual = getattr(multi_layer_encoder, name)
-        desired = module
-        assert actual is desired
+def test_MultiLayerEncoder_trim_registered(mle):
+    mle.register_layer("shallow")
 
-    idx = 1
-    multi_layer_encoder.registered_layers.update([str(idx) for idx in range(idx + 1)])
-    multi_layer_encoder.trim()
+    assert "shallow" in mle
+    assert "intermediate" in mle
+    assert "deep" in mle
 
-    for name, module in modules[: idx + 1]:
-        actual = getattr(multi_layer_encoder, name)
-        desired = module
-        assert actual is desired
+    mle.trim()
 
-    for name in tuple(zip(*modules))[0][idx + 1 :]:
-        with pytest.raises(AttributeError):
-            getattr(multi_layer_encoder, name)
+    assert "shallow" in mle
+    assert "intermediate" not in mle
+    assert "deep" not in mle
 
 
-def test_MultiLayerEncoder_call_future_warning():
-    torch.manual_seed(0)
-    conv = nn.Conv2d(3, 1, 1)
-    input = torch.rand(1, 3, 1, 1)
-
-    modules = (("conv", conv),)
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
-
-    with pytest.warns(FutureWarning):
-        multi_layer_encoder(input, ("conv",))
-
-
-def test_MultiLayerEncoder_encode_future_warning():
-    torch.manual_seed(0)
-    conv = nn.Conv2d(3, 1, 1)
-    input = torch.rand(1, 3, 1, 1)
-
-    modules = (("conv", conv),)
-    multi_layer_encoder = enc.MultiLayerEncoder(modules)
+def test_MultiLayerEncoder_extract_encoder(mle):
+    layer = "intermediate"
+    encoder = mle.extract_encoder(layer)
 
-    with pytest.warns(FutureWarning):
-        multi_layer_encoder.encode(input)
+    assert isinstance(encoder, enc.SingleLayerEncoder)
+    assert encoder.multi_layer_encoder is mle
+    assert encoder.layer == layer
+    assert layer in mle.registered_layers
 
 
-def test_SingleLayerEncoder_call():
+def test_SingleLayerEncoder_call(input):
     torch.manual_seed(0)
     conv = nn.Conv2d(3, 1, 1)
     relu = nn.ReLU(inplace=False)
-    input = torch.rand(1, 3, 128, 128)
 
     modules = (("conv", conv), ("relu", relu))
     multi_layer_encoder = enc.MultiLayerEncoder(modules)
 
-    single_layer_encoder = enc.SingleLayerEncoder(multi_layer_encoder, "conv")
+    single_layer_encoder = enc.SingleLayerEncoder(multi_layer_encoder, "relu")
 
     actual = single_layer_encoder(input)
-    desired = conv(input)
-    ptu.assert_allclose(actual, desired)
+    expected = relu(conv(input))
+    ptu.assert_allclose(actual, expected)
diff --git a/tests/integration/loss/test_multi_op.py b/tests/integration/loss/test_multi_op.py
--- a/tests/integration/loss/test_multi_op.py
+++ b/tests/integration/loss/test_multi_op.py
@@ -76,7 +76,7 @@ def encoder(self):
             return self._encoder
 
         def process_input_image(self, image):
-            return torch.sum(image)
+            return torch.sum(self.encoder(image))
 
     modules = (("count", forward_pass_counter),)
     multi_layer_encoder = enc.MultiLayerEncoder(modules)
diff --git a/tests/integration/test_optim.py b/tests/integration/test_optim.py
--- a/tests/integration/test_optim.py
+++ b/tests/integration/test_optim.py
@@ -1,5 +1,4 @@
 import sys
-from unittest import mock
 
 import pytest
 import pytorch_testing_utils as ptu
@@ -11,6 +10,8 @@
 
 from pystiche import loss, ops, optim
 
+from tests import mocks
+
 skip_if_py38 = pytest.mark.skipif(
     sys.version_info >= (3, 8),
     reason=(
@@ -202,50 +203,7 @@ def image_loader(test_image):
     return data.DataLoader(Dataset(test_image))
 
 
-class ModuleMock(nn.Module):
-    def __init__(self, *methods):
-        super().__init__()
-        self._call_args_list = []
-        for method in methods:
-            setattr(
-                self, method, mock.MagicMock(name=f"{type(self).__name__}.{method}")
-            )
-
-    def forward(self, input, *args, **kwargs):
-        self._call_args_list.insert(0, ((input, *args), kwargs))
-        return input
-
-    @property
-    def call_args_list(self):
-        return self._call_args_list
-
-    @property
-    def called(self):
-        return bool(self.call_args_list)
-
-    @property
-    def call_args(self):
-        return self.call_args_list[0]
-
-    @property
-    def call_count(self):
-        return len(self.call_args_list)
-
-    def assert_called(self):
-        assert self.called
-
-    def assert_called_once(self):
-        assert self.call_count == 1
-
-    def assert_called_once_with(self, input, *args, **kwargs):
-        self.assert_called_once()
-        (input_, *args_), kwargs_ = self.call_args
-        ptu.assert_allclose(input_, input)
-        assert tuple(args_) == args
-        assert kwargs_ == kwargs
-
-
-class TransformerMock(ModuleMock):
+class TransformerMock(mocks.ModuleMock):
     def __init__(self):
         super().__init__()
         self.register_parameter(
@@ -261,12 +219,7 @@ def transformer():
     return TransformerMock()
 
 
-@pytest.fixture
-def module():
-    return ModuleMock()
-
-
-class CriterionMock(ModuleMock):
+class CriterionMock(mocks.ModuleMock):
     def forward(self, *args, **kwargs):
         return torch.sum(super().forward(*args, **kwargs))
 
diff --git a/tests/mocks.py b/tests/mocks.py
--- a/tests/mocks.py
+++ b/tests/mocks.py
@@ -4,6 +4,10 @@
 import unittest.mock
 from distutils import dir_util
 
+import pytorch_testing_utils as ptu
+
+from torch import nn
+
 import pystiche
 
 __all__ = [
@@ -12,6 +16,7 @@
     "ContextMock",
     "patch_models_load_state_dict_from_url",
     "patch_home",
+    "ModuleMock",
 ]
 
 DEFAULT_MOCKER = unittest.mock
@@ -110,3 +115,54 @@ def patch_home(home, copy=True, mocker=DEFAULT_MOCKER):
         dir_util.copy_tree(pystiche.home(), home)
 
     return mocker.patch.dict(os.environ, values={"PYSTICHE_HOME": home})
+
+
+class ModuleMock(nn.Module):
+    def __init__(self, *methods):
+        super().__init__()
+        self._call_args_list = []
+        for method in methods:
+            setattr(
+                self,
+                method,
+                unittest.mock.MagicMock(name=f"{type(self).__name__}.{method}"),
+            )
+
+    def forward(self, input, *args, **kwargs):
+        self._call_args_list.insert(0, ((input, *args), kwargs))
+        return input
+
+    @property
+    def call_args_list(self):
+        return self._call_args_list
+
+    @property
+    def called(self):
+        return bool(self.call_args_list)
+
+    @property
+    def call_args(self):
+        return self.call_args_list[0]
+
+    @property
+    def call_count(self):
+        return len(self.call_args_list)
+
+    def assert_called(self, count=None):
+        assert self.called
+        if count is not None:
+            assert self.call_count == count
+
+    def assert_called_with(self, input, *args, **kwargs):
+        self.assert_called()
+        (input_, *args_), kwargs_ = self.call_args
+        ptu.assert_allclose(input_, input)
+        assert tuple(args_) == args
+        assert kwargs_ == kwargs
+
+    def assert_called_once(self):
+        self.assert_called(count=1)
+
+    def assert_called_once_with(self, input, *args, **kwargs):
+        self.assert_called_once()
+        self.assert_called_with(input, *args, **kwargs)