[WIP] add style transfer task with pystiche

CHANGELOG.md

@@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Refactor preprocess_cls to preprocess, add Serializer, add DataPipelineState ([#229](https://github.com/PyTorchLightning/lightning-flash/pull/229))
 - Added Semantic Segmentation task ([#239](https://github.com/PyTorchLightning/lightning-flash/pull/239) [#287](https://github.com/PyTorchLightning/lightning-flash/pull/287) [#290](https://github.com/PyTorchLightning/lightning-flash/pull/290))
 - Added Object detection prediction example ([#283](https://github.com/PyTorchLightning/lightning-flash/pull/283))
+- Added Style Transfer task and accompanying finetuning and prediction examples ([#262](https://github.com/PyTorchLightning/lightning-flash/pull/262))
 
 ### Changed
 

docs/source/index.rst

@@ -30,6 +30,7 @@ Lightning Flash
    reference/object_detection
    reference/video_classification
    reference/semantic_segmentation
+   reference/style_transfer
 
 .. toctree::
    :maxdepth: 1

docs/source/reference/style_transfer.rst

@@ -0,0 +1,35 @@
+##############
+Style Transfer
+##############
+
+********
+The task
+********
+
+TODO
+
+------
+
+***
+Fit
+***
+
+TODO
+
+------
+
+*************
+API reference
+*************
+
+StyleTransfer
+-------------
+
+.. autoclass:: flash.image.StyleTransfer
+    :members:
+    :exclude-members: forward
+
+StyleTransferData
+-----------------
+
+.. autoclass:: flash.image.StyleTransferData

flash/core/data/data_pipeline.py

@@ -180,7 +180,7 @@ def _resolve_function_hierarchy(
         if object_type is None:
             object_type = Preprocess
 
-        prefixes = ['']
+        prefixes = [None]
         if stage in (RunningStage.TRAINING, RunningStage.TUNING):
             prefixes += ['train', 'fit']
         elif stage == RunningStage.VALIDATING:
@@ -192,7 +192,7 @@ def _resolve_function_hierarchy(
 
         for prefix in prefixes:
             if cls._is_overriden(function_name, process_obj, object_type, prefix=prefix):
-                return f'{prefix}_{function_name}'
+                return function_name if prefix is None else f'{prefix}_{function_name}'
 
         return function_name
 

flash/core/utilities/imports.py

@@ -74,6 +74,7 @@ def _compare_version(package: str, op, version) -> bool:
 _PYTORCHVIDEO_AVAILABLE = _module_available("pytorchvideo")
 _MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
 _TRANSFORMERS_AVAILABLE = _module_available("transformers")
+_PYSTICHE_AVAILABLE = _module_available("pystiche")
 
 if Version:
     _TORCHVISION_GREATER_EQUAL_0_9 = _compare_version("torchvision", operator.ge, "0.9.0")

flash/image/__init__.py

@@ -3,3 +3,4 @@
 from flash.image.detection import ObjectDetectionData, ObjectDetector
 from flash.image.embedding import ImageEmbedder
 from flash.image.segmentation import SemanticSegmentation, SemanticSegmentationData, SemanticSegmentationPreprocess
+from flash.image.style_transfer import StyleTransfer, StyleTransferData, StyleTransferPreprocess

flash/image/style_transfer/__init__.py

@@ -0,0 +1,3 @@
+from flash.image.style_transfer.backbone import STYLE_TRANSFER_BACKBONES
+from flash.image.style_transfer.data import StyleTransferData, StyleTransferPreprocess
+from flash.image.style_transfer.model import StyleTransfer

flash/image/style_transfer/_utils.py

@@ -0,0 +1,10 @@
+from typing import NoReturn
+
+__all__ = ["raise_not_supported"]
+
+
+def raise_not_supported(phase: str) -> NoReturn:
+    raise RuntimeError(
+        f"Style transfer does not support a {phase} phase, "
+        f"since there is no metric to objectively determine the quality of a stylization."
+    )

flash/image/style_transfer/backbone.py

@@ -0,0 +1,23 @@
+import re
+
+from pystiche import enc
+
+from flash.core.registry import FlashRegistry
+
+__all__ = ["STYLE_TRANSFER_BACKBONES"]
+
+MLE_FN_PATTERN = re.compile(r"^(?P<name>\w+?)_multi_layer_encoder$")
+
+STYLE_TRANSFER_BACKBONES = FlashRegistry("backbones")
+
+for mle_fn in dir(enc):
+    match = MLE_FN_PATTERN.match(mle_fn)
+    if not match:
+        continue
+
+    STYLE_TRANSFER_BACKBONES(
+        fn=lambda: (getattr(enc, mle_fn)(), None),
+        name=match.group("name"),
+        namespace="image/style_transfer",
+        package="pystiche",
+    )

flash/image/style_transfer/data.py

@@ -0,0 +1,103 @@
+import functools
+import pathlib
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+import torchvision
+from torch import nn
+from torchvision import transforms
+
+from flash.core.data.data_source import DefaultDataKeys, DefaultDataSources
+from flash.core.data.process import Preprocess
+from flash.core.data.transforms import ApplyToKeys
+from flash.image.classification import ImageClassificationData, ImageClassificationPreprocess
+
+from ._utils import raise_not_supported
+
+__all__ = ["StyleTransferPreprocess", "StyleTransferData"]
+
+
+def _apply_to_input(default_transforms_fn) -> Callable[..., Dict[str, ApplyToKeys]]:
+
+    @functools.wraps(default_transforms_fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Optional[Dict[str, ApplyToKeys]]:
+        default_transforms = default_transforms_fn(*args, **kwargs)
+        if not default_transforms:
+            return default_transforms
+
+        return {hook: ApplyToKeys(DefaultDataKeys.INPUT, transform) for hook, transform in default_transforms.items()}
+
+    return wrapper
+
+
+class StyleTransferPreprocess(ImageClassificationPreprocess):
+
+    def __init__(
+        self,
+        train_transform: Optional[Union[Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Dict[str, Callable]]] = None,
+        predict_transform: Optional[Union[Dict[str, Callable]]] = None,
+        image_size: Union[int, Tuple[int, int]] = 256,
+    ):
+        if val_transform:
+            raise_not_supported("validation")
+        if test_transform:
+            raise_not_supported("test")
+
+        if isinstance(image_size, int):
+            image_size = (image_size, image_size)
+        super().__init__(
+            train_transform=train_transform,
+            val_transform=val_transform,
+            test_transform=test_transform,
+            predict_transform=predict_transform,
+            image_size=image_size,
+        )
+
+    @_apply_to_input
+    def default_transforms(self) -> Optional[Dict[str, Callable]]:
+        if self.training:
+            return dict(
+                to_tensor_transform=torchvision.transforms.ToTensor(),
+                per_sample_transform_on_device=nn.Sequential(
+                    transforms.Resize(min(self.image_size)),
+                    transforms.CenterCrop(self.image_size),
+                ),
+            )
+        elif self.predicting:
+            return dict(
+                to_tensor_transform=torchvision.transforms.ToTensor(),
+                per_sample_transform_on_device=transforms.Resize(min(self.image_size)),
+            )
+        else:
+            # Style transfer doesn't support a validation or test phase, so we return nothing here
+            return None
+
+
+class StyleTransferData(ImageClassificationData):
+    preprocess_cls = StyleTransferPreprocess
+
+    @classmethod
+    def from_folders(
+        cls,
+        train_folder: Optional[Union[str, pathlib.Path]] = None,
+        predict_folder: Optional[Union[str, pathlib.Path]] = None,
+        train_transform: Optional[Union[str, Dict]] = None,
+        predict_transform: Optional[Union[str, Dict]] = None,
+        preprocess: Optional[Preprocess] = None,
+        **kwargs: Any,
+    ) -> "StyleTransferData":
+        if any(param in kwargs for param in ("val_folder", "val_transform")):
+            raise_not_supported("validation")
+        if any(param in kwargs for param in ("test_folder", "test_transform")):
+            raise_not_supported("test")
+
+        preprocess = preprocess or cls.preprocess_cls(train_transform, predict_transform)
+
+        return cls.from_data_source(
+            DefaultDataSources.FOLDERS,
+            train_data=train_folder,
+            predict_data=predict_folder,
+            preprocess=preprocess,
+            **kwargs,
+        )

flash/image/style_transfer/model.py

@@ -0,0 +1,122 @@
+from typing import Any, cast, Dict, Mapping, NoReturn, Optional, Sequence, Type, Union
+
+import pystiche.demo
+import torch
+from pystiche import enc, loss, ops
+from pystiche.image import read_image
+from torch import nn
+from torch.optim.lr_scheduler import _LRScheduler
+
+from flash.core.data.data_source import DefaultDataKeys
+from flash.core.data.process import Serializer
+from flash.core.model import Task
+from flash.core.registry import FlashRegistry
+from flash.image.style_transfer import STYLE_TRANSFER_BACKBONES
+
+from ._utils import raise_not_supported
+
+__all__ = ["StyleTransfer"]
+
+
+class StyleTransfer(Task):
+    backbones: FlashRegistry = STYLE_TRANSFER_BACKBONES
+
+    def __init__(
+        self,
+        style_image: Optional[Union[str, torch.Tensor]] = None,
+        model: Optional[nn.Module] = None,
+        backbone: str = "vgg16",
+        content_layer: str = "relu2_2",
+        content_weight: float = 1e5,
+        style_layers: Sequence[str] = ("relu1_2", "relu2_2", "relu3_3", "relu4_3"),
+        style_weight: float = 1e10,
+        optimizer: Union[Type[torch.optim.Optimizer], torch.optim.Optimizer] = torch.optim.Adam,
+        optimizer_kwargs: Optional[Dict[str, Any]] = None,
+        scheduler: Optional[Union[Type[_LRScheduler], str, _LRScheduler]] = None,
+        scheduler_kwargs: Optional[Dict[str, Any]] = None,
+        learning_rate: float = 1e-3,
+        serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
+    ):
+        self.save_hyperparameters(ignore="style_image")
+
+        if style_image is None:
+            style_image = self.default_style_image()
+        elif isinstance(style_image, str):
+            style_image = read_image(style_image)
+
+        if model is None:
+            model = pystiche.demo.transformer()
+
+        perceptual_loss = self._get_perceptual_loss(
+            backbone=backbone,
+            content_layer=content_layer,
+            content_weight=content_weight,
+            style_layers=style_layers,
+            style_weight=style_weight,
+        )
+        perceptual_loss.set_style_image(style_image)
+
+        super().__init__(
+            model=model,
+            loss_fn=perceptual_loss,
+            optimizer=optimizer,
+            optimizer_kwargs=optimizer_kwargs,
+            scheduler=scheduler,
+            scheduler_kwargs=scheduler_kwargs,
+            learning_rate=learning_rate,
+            serializer=serializer,
+        )
+
+        self.perceptual_loss = perceptual_loss
+
+    def default_style_image(self) -> torch.Tensor:
+        return pystiche.demo.images()["paint"].read(size=256)
+
+    @staticmethod
+    def _modified_gram_loss(encoder: enc.Encoder, *, score_weight: float) -> ops.EncodingComparisonOperator:
+        # The official PyTorch examples as well as the reference implementation of the original author contain an
+        # oversight: they normalize the representation twice by the number of channels. To be compatible with them, we
+        # do the same here.
+        class GramOperator(ops.GramOperator):
+
+            def enc_to_repr(self, enc: torch.Tensor) -> torch.Tensor:
+                repr = super().enc_to_repr(enc)
+                num_channels = repr.size()[1]
+                return repr / num_channels
+
+        return GramOperator(encoder, score_weight=score_weight)
+
+    def _get_perceptual_loss(
+        self,
+        *,
+        backbone: str,
+        content_layer: str,
+        content_weight: float,
+        style_layers: Sequence[str],
+        style_weight: float,
+    ) -> loss.PerceptualLoss:
+        mle, _ = cast(enc.MultiLayerEncoder, self.backbones.get(backbone)())
+        content_loss = ops.FeatureReconstructionOperator(
+            mle.extract_encoder(content_layer), score_weight=content_weight
+        )
+        style_loss = ops.MultiLayerEncodingOperator(
+            mle,
+            style_layers,
+            lambda encoder, layer_weight: self._modified_gram_loss(encoder, score_weight=layer_weight),
+            layer_weights="sum",
+            score_weight=style_weight,
+        )
+        return loss.PerceptualLoss(content_loss, style_loss)
+
+    def training_step(self, batch: Any, batch_idx: int) -> Any:
+        input_image = batch[DefaultDataKeys.INPUT]
+        self.perceptual_loss.set_content_image(input_image)
+
+        output_image = self(input_image)
+        return self.perceptual_loss(output_image).total()
+
+    def validation_step(self, batch: Any, batch_idx: int) -> NoReturn:
+        raise_not_supported("validation")
+
+    def test_step(self, batch: Any, batch_idx: int) -> NoReturn:
+        raise_not_supported("test")

flash_examples/finetuning/style_transfer.py

@@ -0,0 +1,26 @@
+import sys
+
+import flash
+from flash.core.data.utils import download_data
+from flash.core.utilities.imports import _PYSTICHE_AVAILABLE
+
+if _PYSTICHE_AVAILABLE:
+    import pystiche.demo
+else:
+    print("Please, run `pip install pystiche`")
+    sys.exit(0)
+
+from flash.image.style_transfer import StyleTransfer, StyleTransferData
+
+download_data("http://images.cocodataset.org/zips/train2014.zip", "data")
+
+data_module = StyleTransferData.from_folders(train_folder="data", batch_size=4)
+
+style_image = pystiche.demo.images()["paint"].read(size=256)
+
+model = StyleTransfer(style_image)
+
+trainer = flash.Trainer(max_epochs=2)
+trainer.fit(model, data_module)
+
+trainer.save_checkpoint("style_transfer_model.pt")

requirements/datatype_image.txt

@@ -5,3 +5,5 @@ Pillow>=7.2
 kornia>=0.5.1
 matplotlib
 pycocotools>=2.0.2 ; python_version >= "3.7"
+git+https://github.com/pystiche/pystiche@releases/v0.7.2
+# pystiche>=0.7.2

tests/examples/test_scripts.py

@@ -94,6 +94,11 @@ def run_test(filepath):
             "translation.py",
             marks=pytest.mark.skipif(not _TEXT_AVAILABLE, reason="text libraries aren't installed")
         ),
+        # pytest.param(
+        #     "finetuning",
+        #     "style_transfer.py",
+        #     marks=pytest.mark.skipif(not _PYSTICHE_AVAILABLE, reason="pystiche is not installed")
+        # ),  # TODO: takes too long
         pytest.param(
             "predict",
             "image_classification.py",