Changed the loss function to improve training

zhen-he · Aug 2, 2019 · 2f9a239 · 2f9a239
1 parent 81a0766
commit 2f9a239
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/modules/loss_calculator.py b/modules/loss_calculator.py
@@ -22,7 +22,11 @@ def forward(self, output, target, area, **kwargs):
         losses = {}
 
         # Reconstruction loss
-        loss_recon = self.mse(output, target)
+        if 'Y_b' in kwargs.keys():
+            Y_b = kwargs['Y_b']
+            loss_recon = self.mse(output, target) + self.mse((output - Y_b).abs(), (target - Y_b).abs())
+        else：
+            loss_recon = self.mse(output, target)
         loss = loss_recon
         losses['recon'] = loss_recon.data[0]