Switch to Gradient Sensitive Loss

igv · Sep 20, 2018 · 1aa11ab · 1aa11ab
1 parent 467e4e3
commit 1aa11ab
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/ESPCN.py b/ESPCN.py
@@ -52,4 +52,8 @@ def prelu(self, _x, i):
     return tf.nn.relu(_x) - alphas * tf.nn.relu(-_x)
 
   def loss(self, Y, X):
-    return tf.reduce_mean(tf.sqrt(tf.square(X - Y) + 1e-6)) + (1.0 - tf_ssim(Y, X)) * 0.5
+    dY = tf.image.sobel_edges(Y)
+    dX = tf.image.sobel_edges(X)
+    M = tf.sqrt(tf.square(dY[:,:,:,:,0]) + tf.square(dY[:,:,:,:,1]))
+    return tf.losses.absolute_difference(dY, dX) \
+         + tf.losses.absolute_difference((1.0 - M) * Y, (1.0 - M) * X, weights=2.0)
diff --git a/FSRCNN.py b/FSRCNN.py
@@ -85,4 +85,8 @@ def prelu(self, _x, i):
     return tf.nn.relu(_x) - alphas * tf.nn.relu(-_x)
 
   def loss(self, Y, X):
-    return tf.reduce_mean(tf.sqrt(tf.square(X - Y) + 1e-6)) + (1.0 - tf_ssim(Y, X)) * 0.5
+    dY = tf.image.sobel_edges(Y)
+    dX = tf.image.sobel_edges(X)
+    M = tf.sqrt(tf.square(dY[:,:,:,:,0]) + tf.square(dY[:,:,:,:,1]))
+    return tf.losses.absolute_difference(dY, dX) \
+         + tf.losses.absolute_difference((1.0 - M) * Y, (1.0 - M) * X, weights=2.0)