Rewrite TensorFlow L1 norm

Prevent losses from diverging as per tensorflow/tensorflow#12071.
gfabieno · Dec 3, 2020 · b37dbf5 · b37dbf5
1 parent 7754484
commit b37dbf5
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/GeoFlow/Losses.py b/GeoFlow/Losses.py
@@ -63,10 +63,10 @@ def loss(label, output):
 
         # Minimize gradient (blocky inversion)
         if beta > 0:
-            num = tf.norm(output[:, 1:, :] - output[:, :-1, :], ord=1,
+            num = tf_norm(output[:, 1:, :] - output[:, :-1, :],
                           axis=[1, 2])
             if output.get_shape()[-1] != 1:
-                num += tf.norm(output[:, :, 1:] - output[:, :, :-1], ord=1,
+                num += tf_norm(output[:, :, 1:] - output[:, :, :-1],
                                axis=[1, 2])
             den = tf.norm(output, ord=1, axis=[1, 2])
             loss = tf.reduce_sum(num / den)
@@ -75,3 +75,7 @@ def loss(label, output):
         return tf.reduce_sum(losses)
 
     return loss
+
+
+def tf_norm(input, axis=None):
+    return tf.reduce_sum(tf.abs(input), axis=axis)