train_models/mtcnn_model.py: Use OHEM for classification only

Acoording to the original paper, on-line hard example mining technique is used in face/non-face classification task only. Signed-off-by: Ilya Nelkenbaum <[email protected]>
JellyDream · Oct 29, 2017 · 7c1a5ea · 7c1a5ea
1 parent c977fb9
commit 7c1a5ea
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/train_models/mtcnn_model.py b/train_models/mtcnn_model.py
@@ -77,7 +77,8 @@ def bbox_ohem(bbox_pred,bbox_target,label):
     square_error = tf.reduce_sum(square_error,axis=1)
     #keep_num scalar
     num_valid = tf.reduce_sum(valid_inds)
-    keep_num = tf.cast(num_valid*num_keep_radio,dtype=tf.int32)
+    #keep_num = tf.cast(num_valid*num_keep_radio,dtype=tf.int32)
+    keep_num = tf.cast(num_valid, dtype=tf.int32)
     #keep valid index square_error
     square_error = square_error*valid_inds
     _, k_index = tf.nn.top_k(square_error, k=keep_num)
@@ -92,7 +93,8 @@ def landmark_ohem(landmark_pred,landmark_target,label):
     square_error = tf.square(landmark_pred-landmark_target)
     square_error = tf.reduce_sum(square_error,axis=1)
     num_valid = tf.reduce_sum(valid_inds)
-    keep_num = tf.cast(num_valid*num_keep_radio,dtype=tf.int32)
+    #keep_num = tf.cast(num_valid*num_keep_radio,dtype=tf.int32)
+    keep_num = tf.cast(num_valid, dtype=tf.int32)
     square_error = square_error*valid_inds
     _, k_index = tf.nn.top_k(square_error, k=keep_num)
     square_error = tf.gather(square_error, k_index)