Revise V^beta loss

Author: crowsonkb

style_transfer/style_transfer.py

@@ -188,10 +188,6 @@ def v_beta_loss(x, reduction="mean", channel_reduction=None, beta=2.0, eps=1e-8)
     Deep Image Representations by Inverting Them", Mahendran et al (2014)
     (https://arxiv.org/abs/1412.0035).
 
-    The nine-point stencil is from "Fundamental Solutions of 9-point Discrete
-    Laplacians; Derivation and Tables", Lynch (1992)
-    (https://docs.lib.purdue.edu/cgi/viewcontent.cgi?article=1928&context=cstech).
-
     Vectorial total variation was proposed in "Color TV: total variation methods for
     restoration of vector-valued images", Blomgren et al (1998)
     (https://ieeexplore.ieee.org/document/661180).
@@ -213,37 +209,27 @@ def v_beta_loss(x, reduction="mean", channel_reduction=None, beta=2.0, eps=1e-8)
     l, m, h = slice(None, -2), slice(1, -1), slice(2, None)
     x = torch.nn.functional.pad(x, (1, 1, 1, 1), "replicate")
     target = x[..., m, m]
-    ml = (x[..., m, l] - target) ** 2 / 3
-    lm = (x[..., l, m] - target) ** 2 / 3
-    mh = (x[..., m, h] - target) ** 2 / 3
-    hm = (x[..., h, m] - target) ** 2 / 3
-    ll = (x[..., l, l] - target) ** 2 / 12
-    hh = (x[..., h, h] - target) ** 2 / 12
-    hl = (x[..., h, l] - target) ** 2 / 12
-    lh = (x[..., l, h] - target) ** 2 / 12
-    diffs = ml + lm + mh + hm + ll + hh + hl + lh
+    ml = (x[..., m, l] - target) ** 2 / 4  # horizontal 1
+    mh = (x[..., m, h] - target) ** 2 / 4  # horizontal 2
+    lm = (x[..., l, m] - target) ** 2 / 4  # vertical 1
+    hm = (x[..., h, m] - target) ** 2 / 4  # vertical 2
+    ll = (x[..., l, l] - target) ** 2 / 8  # diagonal upper left to lower right 1
+    hh = (x[..., h, h] - target) ** 2 / 8  # diagonal upper left to lower right 2
+    lh = (x[..., l, h] - target) ** 2 / 8  # diagonal lower left to upper right 1
+    hl = (x[..., h, l] - target) ** 2 / 8  # diagonal lower left to upper right 2
+    diffs = ml + mh + lm + hm + ll + hh + lh + hl
     losses = torch.pow(reductions[channel_reduction](diffs, dim=-3) + eps, beta / 2)
     return reductions[reduction](losses)
 
 
 class VBetaLoss(nn.Module):
-    __doc__ = v_beta_loss.__doc__
-
-    def __init__(self, reduction="mean", channel_reduction=None, beta=2.0, eps=1e-8):
+    def __init__(self, beta=2.0, eps=1e-8):
         super().__init__()
-        self.reduction = reduction
-        self.channel_reduction = channel_reduction
         self.beta = beta
         self.eps = eps
 
     def forward(self, x):
-        return v_beta_loss(
-            x,
-            reduction=self.reduction,
-            channel_reduction=self.channel_reduction,
-            beta=self.beta,
-            eps=self.eps,
-        )
+        return v_beta_loss(x * 4, beta=self.beta, eps=self.eps)
 
 
 class SumLoss(nn.ModuleList):
@@ -400,7 +386,7 @@ def get_image(self, image_type='pil'):
     def stylize(self, content_image, style_images, *,
                 style_weights=None,
                 content_weight: float = 0.015,
-                tv_weight: float = 2.,
+                tv_weight: float = 0.125,
                 tv_beta: float = 2.,
                 optimizer: str = 'adam',
                 min_scale: int = 128,