From 3cc33e941411bd0580abad1e58b7d92655824a21 Mon Sep 17 00:00:00 2001
From: haigh1510 <64354653+haigh1510@users.noreply.github.com>
Date: Sat, 9 May 2020 22:02:48 +0500
Subject: [PATCH] RPN: 'compute_loss' fixed

https://github.com/YunYang1994/TensorFlow2.0-Examples/issues/90
---
 4-Object_Detection/RPN/train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/4-Object_Detection/RPN/train.py b/4-Object_Detection/RPN/train.py
index b1f72926..33c2dafa 100644
--- a/4-Object_Detection/RPN/train.py
+++ b/4-Object_Detection/RPN/train.py
@@ -103,14 +103,14 @@ def compute_loss(target_scores, target_bboxes, target_masks, pred_scores, pred_b
     """
     score_loss = tf.nn.softmax_cross_entropy_with_logits(labels=target_scores, logits=pred_scores)
     foreground_background_mask = (np.abs(target_masks) == 1).astype(np.int)
-    score_loss = tf.reduce_sum(score_loss * foreground_background_mask, axis=[1,2,3]) / np.sum(foreground_background_mask)
+    score_loss = tf.reduce_sum(score_loss * foreground_background_mask, axis=[1,2,3]) / np.sum(foreground_background_mask, axis=(1,2,3))
     score_loss = tf.reduce_mean(score_loss)
 
     boxes_loss = tf.abs(target_bboxes - pred_bboxes)
     boxes_loss = 0.5 * tf.pow(boxes_loss, 2) * tf.cast(boxes_loss<1, tf.float32) + (boxes_loss - 0.5) * tf.cast(boxes_loss >=1, tf.float32)
     boxes_loss = tf.reduce_sum(boxes_loss, axis=-1)
     foreground_mask = (target_masks > 0).astype(np.float32)
-    boxes_loss = tf.reduce_sum(boxes_loss * foreground_mask, axis=[1,2,3]) / np.sum(foreground_mask)
+    boxes_loss = tf.reduce_sum(boxes_loss * foreground_mask, axis=[1,2,3]) / np.sum(foreground_mask, axis=(1,2,3))
     boxes_loss = tf.reduce_mean(boxes_loss)
 
     return score_loss, boxes_loss