Repository: pystiche/pystiche

Issue:
Allow supervised datasets in model optimization
Currently we are only allowing unsupervised datasets, i.e. datasets that only return an image during iteration. An example for this would be [`pystiche.data.ImageFolderDataset`](https://github.com/pmeier/pystiche/blob/40dffd99a9d13050faf08702374b6c6d00180c66/pystiche/data/datasets.py#L46-L92). For supervised datasets the user should use [`pystiche.data.Unsupervised`](https://github.com/pmeier/pystiche/blob/40dffd99a9d13050faf08702374b6c6d00180c66/pystiche/data/datasets.py#L18-L29).

This is quite unintuitive and needlessly restrictive. We should simply extract an item from the `image_loader` and use it as `input_image` if it is a `torch.Tensor`. If it is a `tuple` we simply use the first item.


Discussion:

Resolving patch:
diff --git a/pystiche/data/datasets.py b/pystiche/data/datasets.py
--- a/pystiche/data/datasets.py
+++ b/pystiche/data/datasets.py
@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import Any, Callable, Iterator, List, Optional, Tuple, cast
 
 import torch
@@ -8,6 +9,7 @@
 from torchvision.datasets.folder import is_image_file
 
 from pystiche.image import read_image
+from pystiche.misc import build_deprecation_message
 
 __all__ = [
     "Unsupervised",
@@ -17,6 +19,12 @@
 
 class Unsupervised:
     def __init__(self, *args: Any, **kwargs: Any) -> None:
+        msg = build_deprecation_message(
+            "The class 'Unsupervised'",
+            "0.7",
+            info="See https://github.com/pmeier/pystiche/issues/458 for details",
+        )
+        warnings.warn(msg)
         if not isinstance(self, VisionDataset):
             raise RuntimeError
         # This should be used with double inheritance with a
diff --git a/pystiche/optim/optim.py b/pystiche/optim/optim.py
--- a/pystiche/optim/optim.py
+++ b/pystiche/optim/optim.py
@@ -1,7 +1,7 @@
 import time
 import warnings
 from types import TracebackType
-from typing import Any, Callable, Iterable, Optional, Tuple, Type, Union, cast
+from typing import Any, Callable, Iterable, Iterator, Optional, Tuple, Type, Union, cast
 
 import torch
 from torch import nn, optim
@@ -316,6 +316,14 @@ def default_transformer_optimizer(*args: Any, **kwargs: Any) -> Any:
     return default_model_optimizer(*args, **kwargs)
 
 
+def unsupervise(image_loader: Iterable) -> Iterator[torch.Tensor]:
+    for input in image_loader:
+        if isinstance(input, (tuple, list)):
+            input = input[0]
+
+        yield input
+
+
 def model_optimization(
     image_loader: DataLoader,
     transformer: nn.Module,
@@ -333,7 +341,8 @@ def model_optimization(
 
     Args:
         image_loader: Images used as input for the ``transformer``. Drawing from this
-            should yield a single item.
+            should yield either an batched image or a tuple or list with a batched
+            image as first item.
         transformer: Transformer to be optimized.
         criterion: Optimization criterion.
         criterion_update_fn: Is called before each optimization step with the current
@@ -393,7 +402,7 @@ def criterion_update_fn(  # type: ignore[misc]
     )
 
     loading_time_start = time.time()
-    for batch, input_image in enumerate(image_loader, 1):
+    for batch, input_image in enumerate(unsupervise(image_loader), 1):
         input_image = input_image.to(device)
 
         criterion_update_fn(input_image, criterion)  # type: ignore[misc]
@@ -459,7 +468,8 @@ def multi_epoch_model_optimization(
 
     Args:
         image_loader: Images used as input for the ``transformer``. Drawing from this
-            should yield a single item.
+            should yield either an batched image or a tuple or list with a batched
+            image as first item.
         transformer: Transformer to be optimized.
         criterion: Optimization criterion.
         criterion_update_fn: Is called before each optimization step with the current


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
@@ -7,7 +7,7 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from torch.utils.data import DataLoader
+from torch.utils import data
 
 from pystiche import loss, ops, optim
 
@@ -266,9 +266,26 @@ def test_default_transformer_epoch_optim_loop_logging_smoke(caplog, optim_asset_
         )
 
 
+class Dataset(data.Dataset):
+    def __init__(self, image, supervised=False):
+        if image.dim() == 4:
+            image = image.squeeze(0)
+        self.image = image
+        self.supervised = supervised
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, _):
+        if self.supervised:
+            return self.image, torch.zeros(self.image.size()[0])
+        else:
+            return self.image
+
+
 @pytest.fixture
 def image_loader(test_image):
-    return DataLoader((test_image.squeeze(0),))
+    return data.DataLoader(Dataset(test_image))
 
 
 class ModuleMock(nn.Module):
@@ -276,7 +293,13 @@ 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"))
+            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):
@@ -286,15 +309,26 @@ def call_args_list(self):
     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 forward(self, input, *args, **kwargs):
-        self._call_args_list.insert(0, (input, args, kwargs))
-        return input
+    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):
@@ -347,9 +381,9 @@ def calculate_score(
 
 
 def test_model_default_optimization_criterion_update_fn(
-    image_loader, transformer, module, test_image,
+    transformer, test_image,
 ):
-    image_loader = DataLoader((test_image.squeeze(0),))
+    image_loader = data.DataLoader(Dataset(test_image))
 
     content_loss = MSEOperator()
     style_loss = MSEOperator()
@@ -367,3 +401,23 @@ def test_model_optimization_criterion_update_fn_error(
 ):
     with pytest.raises(RuntimeError):
         optim.model_optimization(image_loader, transformer, criterion)
+
+
+@pytest.mark.parametrize(
+    "supervised",
+    (True, False),
+    ids=lambda supervised: f"{'' if supervised else 'un'}supervised",
+)
+def test_model_optimization_image_loader(
+    transformer, criterion, test_image, supervised
+):
+    image_loader = data.DataLoader(Dataset(test_image, supervised=supervised))
+
+    optim.model_optimization(
+        image_loader,
+        transformer,
+        criterion,
+        criterion_update_fn=lambda input_image, criterion: None,
+    )
+
+    transformer.assert_called_once_with(test_image)