Repository: pystiche/pystiche

Issue:
default criterion_update_fn for model optimization
Currently the user has to supply a `criterion_update_fn` to use the model optimization functions. We could set a default if the user supplies a `loss.PerceptualLoss` or `loss.GuidedPerceptualLoss` by invoking `criterion.set_content_image()`.


Discussion:

Resolving patch:
diff --git a/pystiche/optim/optim.py b/pystiche/optim/optim.py
--- a/pystiche/optim/optim.py
+++ b/pystiche/optim/optim.py
@@ -320,7 +320,7 @@ def model_optimization(
     image_loader: DataLoader,
     transformer: nn.Module,
     criterion: nn.Module,
-    criterion_update_fn: Callable[[torch.Tensor, nn.Module], None],
+    criterion_update_fn: Optional[Callable[[torch.Tensor, nn.Module], None]] = None,
     optimizer: Optional[Optimizer] = None,
     get_optimizer: Optional[Callable[[nn.Module], Optimizer]] = None,
     quiet: bool = False,
@@ -337,7 +337,10 @@ def model_optimization(
         transformer: Transformer to be optimized.
         criterion: Optimization criterion.
         criterion_update_fn: Is called before each optimization step with the current
-            images and the optimization ``criterion``.
+            images and the optimization ``criterion``. If omitted and ``criterion`` is
+            a :class:`~pystiche.loss.PerceptualLoss` or a
+            :class:`~pystiche.loss.GuidedPerceptualLoss` this defaults to invoking
+            :meth:`~pystiche.loss.PerceptualLoss.set_content_image`.
         optimizer: Optional optimizer. If ``None``,
             :func:`default_model_optimizer` is used.
         quiet: If ``True``, not information is logged during the optimization. Defaults
@@ -350,6 +353,22 @@ def model_optimization(
             :func:`pystiche.optim.default_transformer_optim_log_fn` is used. Defaults
             to ``None``.
     """
+    if criterion_update_fn is None:
+        if isinstance(criterion, (loss.PerceptualLoss, loss.GuidedPerceptualLoss)):
+
+            def criterion_update_fn(  # type: ignore[misc]
+                input_image: torch.Tensor,
+                criterion: Union[loss.PerceptualLoss, loss.GuidedPerceptualLoss],
+            ) -> None:
+                criterion.set_content_image(input_image)
+
+        else:
+            raise RuntimeError(
+                f"The parameter 'criterion_update_fn' can only be omitted if the "
+                f"'criterion' is a loss.PerceptualLoss or a loss.GuidedPerceptualLoss. "
+                f"Got {type(criterion)} instead."
+            )
+
     if optimizer is None:
         optimizer = default_model_optimizer(transformer)
 
@@ -377,7 +396,7 @@ def model_optimization(
     for batch, input_image in enumerate(image_loader, 1):
         input_image = input_image.to(device)
 
-        criterion_update_fn(input_image, criterion)
+        criterion_update_fn(input_image, criterion)  # type: ignore[misc]
 
         loading_time = time.time() - loading_time_start
 
@@ -423,8 +442,8 @@ def multi_epoch_model_optimization(
     image_loader: DataLoader,
     transformer: nn.Module,
     criterion: nn.Module,
-    criterion_update_fn: Callable[[torch.Tensor, nn.Module], None],
-    epochs: int,
+    criterion_update_fn: Optional[Callable[[torch.Tensor, nn.Module], None]] = None,
+    epochs: int = 2,
     optimizer: Optional[Optimizer] = None,
     lr_scheduler: Optional[LRScheduler] = None,
     quiet: bool = False,
@@ -444,8 +463,11 @@ def multi_epoch_model_optimization(
         transformer: Transformer to be optimized.
         criterion: Optimization criterion.
         criterion_update_fn: Is called before each optimization step with the current
-            images and the optimization ``criterion``.
-        epochs: Number of epochs.
+            images and the optimization ``criterion``. If omitted and ``criterion`` is
+            a :class:`~pystiche.loss.PerceptualLoss` or a
+            :class:`~pystiche.loss.GuidedPerceptualLoss` this defaults to invoking
+            :meth:`~pystiche.loss.PerceptualLoss.set_content_image`.
+        epochs: Number of epochs. Defaults to ``2``.
         optimizer: Optional optimizer. If ``None``, it is extracted from
             ``lr_scheduler`` or func:`default_model_optimizer` is used.
         lr_scheduler: Optional learning rate scheduler. ``step()`` is invoked after


Test patch:
diff --git a/tests/integration/optim/test_optim.py b/tests/integration/optim/test_optim.py
--- a/tests/integration/optim/test_optim.py
+++ b/tests/integration/optim/test_optim.py
@@ -1,12 +1,15 @@
 import sys
+from unittest import mock
 
 import pytest
 import pytorch_testing_utils as ptu
 
 import torch
+import torch.nn.functional as F
 from torch import nn
+from torch.utils.data import DataLoader
 
-from pystiche import optim
+from pystiche import loss, ops, optim
 
 from tests import asserts
 
@@ -261,3 +264,106 @@ def test_default_transformer_epoch_optim_loop_logging_smoke(caplog, optim_asset_
             logger=optim_logger,
             log_fn=log_fn,
         )
+
+
+@pytest.fixture
+def image_loader(test_image):
+    return DataLoader((test_image.squeeze(0),))
+
+
+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"))
+
+    @property
+    def call_args_list(self):
+        return self._call_args_list
+
+    @property
+    def called(self):
+        return bool(self.call_args_list)
+
+    def call_args(self):
+        return self.call_args_list[0]
+
+    def call_count(self):
+        return len(self.call_args_list)
+
+    def forward(self, input, *args, **kwargs):
+        self._call_args_list.insert(0, (input, args, kwargs))
+        return input
+
+
+class TransformerMock(ModuleMock):
+    def __init__(self):
+        super().__init__()
+        self.register_parameter(
+            "parameter", nn.Parameter(torch.zeros(1).squeeze(), requires_grad=True)
+        )
+
+    def forward(self, input, *args, **kwargs):
+        return super().forward(input + self.parameter, *args, **kwargs)
+
+
+@pytest.fixture
+def transformer():
+    return TransformerMock()
+
+
+@pytest.fixture
+def module():
+    return ModuleMock()
+
+
+class CriterionMock(ModuleMock):
+    def forward(self, *args, **kwargs):
+        return torch.sum(super().forward(*args, **kwargs))
+
+
+@pytest.fixture
+def criterion():
+    return CriterionMock()
+
+
+class MSEOperator(ops.PixelComparisonOperator):
+    def image_to_repr(self, image):
+        return image
+
+    def input_image_to_repr(
+        self, image, ctx,
+    ):
+        return image
+
+    def target_image_to_repr(self, image):
+        return image, None
+
+    def calculate_score(
+        self, input_repr, target_repr, ctx,
+    ):
+        return F.mse_loss(input_repr, target_repr)
+
+
+def test_model_default_optimization_criterion_update_fn(
+    image_loader, transformer, module, test_image,
+):
+    image_loader = DataLoader((test_image.squeeze(0),))
+
+    content_loss = MSEOperator()
+    style_loss = MSEOperator()
+    criterion = loss.PerceptualLoss(content_loss, style_loss)
+
+    content_loss.set_target_image(torch.rand_like(test_image))
+    style_loss.set_target_image(torch.rand_like(test_image))
+    optim.model_optimization(image_loader, transformer, criterion)
+
+    ptu.assert_allclose(content_loss.target_image, test_image)
+
+
+def test_model_optimization_criterion_update_fn_error(
+    image_loader, transformer, criterion
+):
+    with pytest.raises(RuntimeError):
+        optim.model_optimization(image_loader, transformer, criterion)