Add a safe wrapper for cross entropy.

luminoth/models/fasterrcnn/rcnn.py

@@ -5,6 +5,9 @@
 from luminoth.models.fasterrcnn.rcnn_target import RCNNTarget
 from luminoth.models.fasterrcnn.roi_pool import ROIPoolingLayer
 from luminoth.utils.losses import smooth_l1_loss
+from luminoth.utils.safe_wrappers import (
+    safe_softmax_cross_entropy_with_logits
+)
 from luminoth.utils.vars import (
     get_initializer, layer_summaries, variable_summaries,
     get_activation_function
@@ -304,7 +307,7 @@ def loss(self, prediction_dict):
 
             # We get cross entropy loss of each proposal.
             cross_entropy_per_proposal = (
-                tf.nn.softmax_cross_entropy_with_logits(
+                safe_softmax_cross_entropy_with_logits(
                     labels=cls_target_one_hot, logits=cls_score_labeled
                 )
             )

luminoth/models/fasterrcnn/rpn.py

@@ -10,6 +10,9 @@
 from .rpn_target import RPNTarget
 from .rpn_proposal import RPNProposal
 from luminoth.utils.losses import smooth_l1_loss
+from luminoth.utils.safe_wrappers import (
+    safe_softmax_cross_entropy_with_logits
+)
 from luminoth.utils.vars import (
     get_initializer, layer_summaries, variable_summaries,
     get_activation_function
@@ -257,7 +260,7 @@ def loss(self, prediction_dict):
             cls_target = tf.one_hot(labels, depth=2)
 
             # Equivalent to log loss
-            ce_per_anchor = tf.nn.softmax_cross_entropy_with_logits(
+            ce_per_anchor = safe_softmax_cross_entropy_with_logits(
                 labels=cls_target, logits=cls_score
             )
             prediction_dict['cross_entropy_per_anchor'] = ce_per_anchor

luminoth/utils/safe_wrappers.py

@@ -0,0 +1,16 @@
+import tensorflow as tf
+
+
+def safe_softmax_cross_entropy_with_logits(
+   labels, logits, name='safe_cross_entropy'):
+    with tf.name_scope(name):
+        safety_condition = tf.greater(
+            tf.shape(logits)[0], 0, name='safety_condition'
+        )
+        return tf.cond(
+            safety_condition,
+            true_fn=lambda: tf.nn.softmax_cross_entropy_with_logits(
+                labels=labels, logits=logits
+            ),
+            false_fn=lambda: tf.constant([], dtype=logits.dtype)
+        )