Added StyleLoss module

Author: myscience

src/losses.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+from torchvision.models import get_model
+from torch.nn.functional import mse_loss
+
+from typing import List
+from torch import Tensor
+
+from .recorder import Traced
+from .utils import compute_stat
+
+def style_loss(targ : Tensor, pred : Tensor, k : float = 0.) -> Tensor:
+    bs, c, *_ = targ.shape
+
+    # Compare mean & std of target and pred tensor
+    targ_mean, targ_std = compute_stat(targ)
+    pred_mean, pred_std = compute_stat(pred)
+
+    loss_mean = mse_loss(targ_mean, pred_mean, reduction='none')
+    loss_std  = mse_loss(targ_std , pred_std , reduction='none')
+
+    if k > 0:
+        loss_mean, sort_idx = torch.sort(loss_mean, dim=1)
+
+        loss_mean[:, int(c * k):] = 0
+        loss_std[:, sort_idx[:, int(c * k):]] = 0
+
+    return (loss_mean + loss_std).mean()
+
+def content_loss(targ : Tensor, pred : Tensor) -> Tensor:
+    # Compare mean & std of target and pred tensor
+    targ_mean, targ_std = compute_stat(targ)
+    pred_mean, pred_std = compute_stat(pred)
+
+    norm_targ = (targ - targ_mean) / targ_std
+    norm_pred = (pred - pred_mean) / pred_std
+
+    return mse_loss(norm_targ, norm_pred)
+
+class StyleLoss(nn.Module):
+    '''
+        This module implements the style loss used for model
+        training in the paper:
+        `Hierarchy Flow For High-Fidelity Image-to-Image Translation`
+        Fan et al. (2023) (arxiv:2308.06909).
+
+        This loss is a combination of a standard VGG-19 (style) loss
+        and and an original (modification of) content loss. The tradeoff
+        between content-preserving and style-preserving is controlled by
+        the `content_weight` parameter.
+
+        Args:
+            content_weight (float): The weight of the content loss.
+            style_weight (float): The weight of the style loss.
+            vgg_layers (list): The layers of the VGG-19 model to use. 
+    '''
+
+    def __init__(
+        self,
+        enc_depth : List[int] = (3, 10, 17, 30),
+        backbone : str = 'vgg19_bn',
+        content_weight : float = .8,
+    ) -> None:
+        super().__init__()
+
+        backbone = get_model(backbone, weights='DEFAULT')
+        backbone = Traced(backbone, enc_depth)
+
+        self.backbone = backbone
+
+        self.content_weight = content_weight
+
+    def forward(
+        self,
+        orig_img : Tensor,
+        targ_sty : Tensor,
+        pred_img : Tensor,
+        align_k : float = 0.8    
+    ) -> Tensor:
+        '''
+        '''
+        # Get features from the backbone model for the
+        # original content image
+        (*_, orig_feat) = self.backbone(orig_img)
+
+        # Get the style features for both the reference
+        # and produced images
+        targ_feats = self.backbone(targ_sty)
+        pred_feats = self.backbone(pred_img)
+
+        # Compute the style & content loss
+        loss_style = sum([style_loss(f1, f2, k=align_k) for f1, f2 in zip(targ_feats, pred_feats)])
+        loss_content = content_loss(orig_feat, pred_feats[-1])
+
+        # Combine the style and content loss
+        loss = loss_content + self.content_weight * loss_style
+
+        return loss
+
+

src/recorder.py

@@ -0,0 +1,93 @@
+
+import torch.nn as nn
+from torch import Tensor
+
+from typing import List
+
+from .utils import default
+from .utils import flatten
+
+class FeatureRecorder:
+    ''' 
+        Basic feature recording class that implements a PyTorch hook to
+        acquire a layer activations as they get processed by the network
+    '''
+    
+    def __init__(
+        self,
+        names : List[str],
+    ) -> None:
+        
+        self.names = names
+        self.feats = {l : [] for l in names}
+        
+    def __call__(self, module : nn.Module, inp : Tensor, out : Tensor) -> None:
+        # Detach layer output from PyTorch graph
+        data = out.detach()
+
+        # Get the module name
+        layer = module.name
+
+        self.feats[layer].append(data)
+
+    def clean(
+        self,
+        names : List[str] | None = None,
+    ) -> None:
+        self.names = default(names, self.names)
+        self.feats = {k : [] for k in self.names}
+
+class Traced(nn.Module):
+    '''
+        A wrapper class of a Torch Module whose intermediate activations
+        are traced (recorded) via forward hooks at chosen depth.
+    '''
+
+    def __init__(
+        self,
+        module : nn.Module,
+        depths : List[int],
+    ) -> None:
+        super().__init__()
+
+        self.module = module
+        self.depths = depths
+
+        # Get the list of layers to be traced
+        self.layers = [l for depth, l in enumerate(flatten(module)) if depth in depths]
+
+        for layer, name in zip(self.layers, self.names): layer.name = name
+
+        # Initialize the feature recorder
+        self.names = [f'enc_{d}' for d in depths]
+        self.recorder = FeatureRecorder(self.names)
+
+        # Register the forward hooks for each layered targeted as 'traced'
+        self.hook_handles = [l.register_forward_hook(self.recorder) for l in self.layers]
+    
+    @property
+    def features(self) -> List[Tensor]:
+        return [self.recorder.feats[k] for k in self.names]
+
+    def forward(self, inp : Tensor, auto_clean : bool = True) -> List[Tensor]:
+        # Propagate the input into the network
+        _ = self.module(inp)
+
+        # Collect the features
+        feats = self.features
+
+        if auto_clean: self.clean()
+
+        return feats
+
+    def clean(self) -> None:
+        self.recorder.clean()
+
+        for h in self.hook_handles: h.remove()
+
+    def __str__(self) -> str:
+        msg = 'Tracing module: \n'
+        msg += f'{self.module} \n'
+        msg += f'Traced layers: {self.layers} \n'
+        msg += f'Traced depths: {self.depths} \n'
+        return msg

src/utils.py

@@ -2,13 +2,36 @@
 import torch.nn as nn
 
 from torch import Tensor
-from typing import Any, Tuple
+from typing import Any, Tuple, List
 
 from einops import rearrange
 
+def exists(var : Any | None) -> bool:
+    return var is not None
+
 def default(var : Any | None, val : Any) -> Any:
     return val if var is None else var
 
+def flatten(model : nn.Module, exclude : List[nn.Module] = []) -> List[nn.Module]:
+    flattened = [flatten(children) for children in model.children()]
+    res = [model] if list(model.children()) == [] else []
+
+    for c in flattened: res += c
+    
+    return res
+
+
+def compute_stat(feat : Tensor, eps : float = 1e-5) -> Tuple[Tensor, Tensor]:
+    # Check input dimension
+    bs, c, h, w = feat.shape
+
+    var = rearrange(feat, 'b c h w -> b c (h w)').var(dim=2) + eps
+    std = rearrange(var.sqrt(), 'b c -> b c 1 1')
+
+    mean = rearrange(feat, 'b c h w -> b c (h w)').mean(dim=2)
+    mean = rearrange(mean, 'b c -> b c 1 1')
+
+    return mean, std
 
 class ReversibleConcat(nn.Module):
     '''
@@ -59,20 +82,8 @@ def __init__(self) -> None:
 
     def forward(self, subj : Tensor, feat_mean : Tensor, feat_std : Tensor) -> Tensor:
         # Get subject mean and standard deviation
-        subj_mean, subj_std = self._compute_stat(subj)
+        subj_mean, subj_std = compute_stat(subj)
 
         norm_feat = (subj - subj_mean) / subj_std
 
-        return norm_feat * feat_std + feat_mean
-
-    def _compute_stat(self, feat : Tensor, eps : float = 1e-5) -> Tuple[Tensor, Tensor]:
-        # Check input dimension
-        bs, c, h, w = feat.shape
-
-        var = rearrange(feat, 'b c h w -> b c (h w)').var(dim=2) + eps
-        std = rearrange(var.sqrt(), 'b c -> b c 1 1')
-
-        mean = rearrange(feat, 'b c h w -> b c (h w)').mean(dim=2)
-        mean = rearrange(mean, 'b c -> b c 1 1')
-
-        return mean, std
+        return norm_feat * feat_std + feat_mean