Fixed bug in loss, added backward test

Author: myscience

src/losses.py

@@ -58,7 +58,7 @@ class StyleLoss(nn.Module):
 
     def __init__(
         self,
-        enc_depth : List[int] = (3, 10, 17, 30),
+        enc_depth : List[int] = (5, 12, 19, 32),
         backbone : str = 'vgg19_bn',
         content_weight : float = .8,
     ) -> None:

src/recorder.py

@@ -23,13 +23,11 @@ def __init__(
         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)
+        # NOTE: It's important NOT to detach output from grad graph
+        self.feats[layer].append(out)
 
     def clean(
         self,

test/test_hflow.py

@@ -1,7 +1,10 @@
 import torch
 import unittest
 
+from torch import Tensor
+
 from src.hflow import HierarchyFlow
+from src.losses import StyleLoss
 
 class HFlowTest(unittest.TestCase):
     def setUp(self) -> None:
@@ -19,6 +22,11 @@ def setUp(self) -> None:
             ]
         )
 
+        self.criterion = StyleLoss(
+            enc_depth=(5, 12, 19, 32),
+            backbone='vgg19_bn',
+        )
+
         # Create a dummy input of correct shape
         self.input_shape = (2, 3, 256, 256)
 
@@ -30,3 +38,14 @@ def test_forward(self):
 
         # Check the output shape
         self.assertEqual(output.shape, self.input_shape)
+
+    def test_loss_backward(self):
+        pred_img = self.flow(self.dummy_subj, self.dummy_style)
+        
+        loss : Tensor = self.criterion(self.dummy_subj, self.dummy_style, pred_img)
+
+        loss.backward()
+
+        # Check the output shape
+        self.assertTrue(loss > 0)
+

test/test_loss.py

@@ -1,25 +1,42 @@
 import torch
 import unittest
 
+import torch.nn as nn
+from torch import Tensor
+
 from src.losses import StyleLoss
 
 class StyleLossTest(unittest.TestCase):
     def setUp(self) -> None:
 
         self.criterion = StyleLoss(
-            enc_depth=(3, 10, 17, 30),
+            enc_depth=(5, 12, 19, 32),
             backbone='vgg19_bn',
         )
 
         # Create a dummy input of correct shape
         self.input_shape = (2, 3, 256, 256)
 
-        self.orig_img = torch.randn(*self.input_shape)
-        self.targ_sty = torch.randn(*self.input_shape)
-        self.pred_img = torch.randn(*self.input_shape)
+        self.orig_img = torch.randn(*self.input_shape, requires_grad=True)
+        self.targ_sty = torch.randn(*self.input_shape, requires_grad=True)
+        self.pred_img = torch.randn(*self.input_shape, requires_grad=True)
+
+    def test_default_is_relu(self):
+        layers = self.criterion.backbone.layers
 
+        self.assertTrue(all([isinstance(l, nn.ReLU) for l in layers]))
+
     def test_forward(self):
         loss = self.criterion(self.orig_img, self.targ_sty, self.pred_img)
 
         # Check the output shape
         self.assertTrue(loss > 0)
+
+    def test_backward(self):
+        loss : Tensor = self.criterion(self.orig_img, self.targ_sty, self.pred_img)
+
+        loss.backward()
+
+        # Check the output shape
+        self.assertTrue(loss > 0)
+