Added logging for loss debugging

1508.06576/loss.py

@@ -56,6 +56,13 @@ class Loss:
         @params
         Author: @jiteshg
         '''
+        loss_t = 0.0
+        a = 0.0
+        b = 0.0
         for gen,cont,sty in zip(content_fmap_gen,cont_fmap_real,style_fmap_real):
-            loss_t = alpha*Loss.content_loss(cont,gen) + beta*Loss.style_loss(sty,gen)
-        return loss_t
\ No newline at end of file
+            loss_cont = Loss.content_loss(cont,gen)
+            loss_style = Loss.style_loss(sty,gen)
+            a+= loss_cont
+            b+= loss_style
+            loss_t += alpha*loss_cont + beta*loss_style
+        return loss_t,a,b
\ No newline at end of file

1508.06576/optimizer.py

@@ -17,9 +17,9 @@ class Optimizer:
         '''
         LOGGER.info("Running gradient descent with the following parameters")
         epoch = 5000
-        learning_rate = 0.001
-        alpha = 10
-        beta = 100
+        learning_rate = 0.002
+        alpha = 1
+        beta = 0.01 
         LOGGER.info(f"{epoch},{learning_rate},{alpha},{beta}")
 
         optimizer=optim.Adam([content_img_clone],lr=learning_rate)
@@ -33,7 +33,7 @@ class Optimizer:
             style_fmaps = feature_maps.get_fmaps(style_img)
             content_generated_fmaps = feature_maps.get_fmaps(content_img_clone)
             
-            total_loss = Loss.total_loss(alpha, beta, content_fmaps, style_fmaps, content_generated_fmaps)
+            total_loss,total_cont_loss,total_style_loss = Loss.total_loss(alpha, beta, content_fmaps, style_fmaps, content_generated_fmaps)
             
             # clears x.grad for every parameter x in the optimizer. 
             # It’s important to call this before total_loss.backward(), otherwise it will accumulate the gradients from multiple passes.
@@ -46,6 +46,6 @@ class Optimizer:
             optimizer.step()
             #plt.clf()
             #plt.plot(content_img_clone)
-            if(e%10):
-                LOGGER.info(f"Epoch = {e} Total Loss = {total_loss}")
+            if(e%10 == 0):
+                LOGGER.info(f"Epoch = {e} Total Loss = {total_loss} Style Loss = {total_cont_loss} Content Loss = {total_style_loss}")
                 save_image(content_img_clone,"styled.png")
\ No newline at end of file

1508.06576/preprocess.py

@@ -35,9 +35,10 @@ class Preprocessor:
         @params
         img: 3d numpy array
         '''
-        loader = transforms.Compose([transforms.ToTensor(),transforms.Resize([512,512]),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225],),])
+        #loader = transforms.Compose([transforms.ToTensor(),transforms.Resize([224,224]),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225],),])
+        loader = transforms.Compose([transforms.ToTensor(),transforms.Resize([224,224])])
         img = loader(img).unsqueeze(0)
-        assert img.shape == (1,3,512,512)
+        assert img.shape == (1,3,224,224)
         return img.to(device,torch.float)