core\refac: #98 debug and fix loss

- debug and fix tgce loss variants

Author: blotero

core/seg_tgce/experiments/histology/features.py

@@ -13,15 +13,18 @@
 
 from ..utils import handle_training
 
-TARGET_SHAPE = (256, 256)
-BATCH_SIZE = 32
+TARGET_SHAPE = (128, 128)
+BATCH_SIZE = 4
 TRAIN_EPOCHS = 20
 TUNER_EPOCHS = 1
 
 DEFAULT_HPARAMS = {
     "initial_learning_rate": 1e-3,
     "q": 0.5,
     "noise_tolerance": 0.5,
+    "a": 0.5,
+    "b": 0.5,
+    "c": 1.0,
     "lambda_reg_weight": 0.1,
     "lambda_entropy_weight": 0.1,
     "lambda_sum_weight": 0.1,
@@ -49,12 +52,18 @@ def build_model(hp=None):
             "lambda_sum_weight": hp.Float(
                 "lambda_sum_weight", min_value=0.01, max_value=0.5, step=0.01
             ),
+            "a": hp.Float("a", min_value=0.0, max_value=1.0, step=0.1),
+            "b": hp.Float("b", min_value=0.0, max_value=1.0, step=0.1),
+            "c": hp.Float("c", min_value=0.0, max_value=1.0, step=0.1),
         }
 
     return build_features_model_from_hparams(
         learning_rate=params["initial_learning_rate"],
         q=params["q"],
         noise_tolerance=params["noise_tolerance"],
+        a=params["a"],
+        b=params["b"],
+        c=params["c"],
         lambda_reg_weight=params["lambda_reg_weight"],
         lambda_entropy_weight=params["lambda_entropy_weight"],
         lambda_sum_weight=params["lambda_sum_weight"],
@@ -76,10 +85,7 @@ def build_model(hp=None):
     args = parser.parse_args()
 
     processed_train, processed_validation, processed_test = get_processed_data(
-        image_size=TARGET_SHAPE,
-        batch_size=BATCH_SIZE,
-        use_augmentation=True,
-        augmentation_factor=2,
+        image_size=TARGET_SHAPE, batch_size=BATCH_SIZE, use_augmentation=False
     )
 
     model = handle_training(

core/seg_tgce/loss/test_tgce.py

@@ -0,0 +1,43 @@
+import tensorflow as tf
+
+from seg_tgce.data.oxford_pet.oxford_pet_tfds import get_data_multiple_annotators_tfds
+from seg_tgce.loss.tgce import TgceFeatures, TgcePixel, TgceScalar
+
+if __name__ == "__main__":
+    NOISE_LEVELS_SNR = [-20.0, 20.0]
+    TARGET_SHAPE = (128, 128)
+    BATCH_SIZE = 16
+    NUM_CLASSES = 3
+
+    train_data, val_data, test_data = get_data_multiple_annotators_tfds(
+        noise_levels_snr=NOISE_LEVELS_SNR,
+        target_shape=TARGET_SHAPE,
+        batch_size=BATCH_SIZE,
+        labeling_rate=0.7,
+    )
+    scalar_loss = TgceScalar(num_classes=NUM_CLASSES)
+    features_loss = TgceFeatures(num_classes=NUM_CLASSES)
+    pixel_loss = TgcePixel(num_classes=NUM_CLASSES)
+
+    for i, data in enumerate(train_data):
+        ground_truth = data["ground_truth"]
+        images = data["image"]
+        masks = data["masks"]
+        labeler_mask = data["labeler_mask"]
+        print(ground_truth.shape)
+        print(images.shape)
+        print(masks.shape)
+        print(labeler_mask.shape)
+        # lambda r with zeroes
+        scalar_lambda_r = tf.ones((BATCH_SIZE, len(NOISE_LEVELS_SNR)))
+        features_lambda_r = tf.ones((BATCH_SIZE, 4, 4, len(NOISE_LEVELS_SNR)))
+        pixel_lambda_r = tf.ones((BATCH_SIZE, *TARGET_SHAPE, len(NOISE_LEVELS_SNR)))
+        y_pred = tf.random.normal((BATCH_SIZE, *TARGET_SHAPE, NUM_CLASSES))
+
+        scalar_loss_value = scalar_loss.call(masks, y_pred, scalar_lambda_r, labeler_mask)
+        features_loss_value = features_loss.call(
+            masks, y_pred, features_lambda_r, labeler_mask
+        )
+        pixel_loss_value = pixel_loss.call(masks, y_pred, pixel_lambda_r, labeler_mask)
+
+        break

core/seg_tgce/loss/tgce.py

@@ -188,7 +188,11 @@ def call(
         y_pred = tf.clip_by_value(y_pred, self.epsilon, 1.0 - self.epsilon)
         lambda_r = tf.clip_by_value(lambda_r, self.epsilon, 1.0 - self.epsilon)
 
-        reg_term = reliability_penalizer(labeler_mask, lambda_r, self.a, self.b, self.c)
+        lambda_r_reduced = tf.reduce_mean(lambda_r, axis=(1, 2))
+
+        reg_term = reliability_penalizer(
+            labeler_mask, lambda_r_reduced, self.a, self.b, self.c
+        )
         # Expand predictions to match annotators dimension
         y_pred_exp = tf.expand_dims(y_pred, axis=-1)
         y_pred_exp = tf.tile(y_pred_exp, [1, 1, 1, 1, tf.shape(y_true)[-1]])
@@ -222,7 +226,7 @@ def call(
         )
 
         lambda_sum = self.lambda_sum_weight * tf.reduce_mean(
-            tf.square(tf.reduce_sum(valid_lambda_r, axis=-2) - 1.0)
+            tf.square(tf.reduce_sum(valid_lambda_r, axis=-1) - 1.0)
         )
 
         total_loss = (
@@ -305,7 +309,12 @@ def call(
 
         y_pred = tf.clip_by_value(y_pred, self.epsilon, 1.0 - self.epsilon)
         lambda_r = tf.clip_by_value(lambda_r, self.epsilon, 1.0 - self.epsilon)
-        reg_term = reliability_penalizer(labeler_mask, lambda_r, self.a, self.b, self.c)
+
+        lambda_r_reduced = tf.reduce_mean(lambda_r, axis=(1, 2))
+
+        reg_term = reliability_penalizer(
+            labeler_mask, lambda_r_reduced, self.a, self.b, self.c
+        )
 
         # Expand predictions to match annotators dimension
         y_pred_exp = tf.expand_dims(y_pred, axis=-1)