From be77f035c85de4db93ac3c50e42fc85e931cef35 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Fri, 13 Jan 2023 17:58:00 +0000 Subject: [PATCH] dlr: add cross-entropy + dice loss fn option --- aimodel/slurm-TEST-deeplabv3p-rainfall.job | 5 ++- aimodel/src/deeplabv3_plus_test_rainfall.py | 19 +++++++-- .../lib/ai/components/LossCrossEntropyDice.py | 41 +++++++++++++++++++ 3 files changed, 59 insertions(+), 6 deletions(-) create mode 100644 aimodel/src/lib/ai/components/LossCrossEntropyDice.py diff --git a/aimodel/slurm-TEST-deeplabv3p-rainfall.job b/aimodel/slurm-TEST-deeplabv3p-rainfall.job index 2621765..0febcfe 100755 --- a/aimodel/slurm-TEST-deeplabv3p-rainfall.job +++ b/aimodel/slurm-TEST-deeplabv3p-rainfall.job @@ -30,10 +30,11 @@ show_help() { echo -e " DIR_RAINFALLWATER The path to the directory the .tfrecord files containing the rainfall radar / water depth data." >&2; echo -e " PATH_HEIGHTMAP The path to the heightmap jsonl file to read in." >&2; echo -e " PATH_COLOURMAP The path to the colourmap for predictive purposes." >&2; - echo -e " PATH_CHECKPOINT The path to a checkcpoint to load. If specified, a model will be loaded instead of being trained." >&2; echo -e " STEPS_PER_EPOCH The number of steps to consider an epoch. Defaults to None, which means use the entire dataset." >&2; echo -e " NO_REMOVE_ISOLATED_PIXELS Set to any value to avoid the engine from removing isolated pixels - that is, water pixels with no other surrounding pixels, either side to side to diagonally." >&2; echo -e " EPOCHS The number of epochs to train for." >&2; + echo -e " LOSS The loss function to use. Default: cross-entropy (possible values: cross-entropy, cross-entropy-dice)." >&2; + echo -e " PATH_CHECKPOINT The path to a checkcpoint to load. If specified, a model will be loaded instead of being trained." >&2; echo -e " PREDICT_COUNT The number of items from the (SCRAMBLED) dataset to make a prediction for." >&2; echo -e " POSTFIX Postfix to append to the output dir (auto calculated)." >&2; echo -e " ARGS Optional. Any additional arguments to pass to the python program." >&2; @@ -60,7 +61,7 @@ DIR_OUTPUT="output/$(date -u --rfc-3339=date)_${CODE}"; echo -e ">>> Additional args: ${ARGS}"; export PATH=$HOME/software/bin:$PATH; -export IMAGE_SIZE BATCH_SIZE DIR_RAINFALLWATER PATH_HEIGHTMAP PATH_COLOURMAP STEPS_PER_EPOCH DIR_OUTPUT PATH_CHECKPOINT EPOCHS PREDICT_COUNT NO_REMOVE_ISOLATED_PIXELS; +export IMAGE_SIZE BATCH_SIZE DIR_RAINFALLWATER PATH_HEIGHTMAP PATH_COLOURMAP STEPS_PER_EPOCH DIR_OUTPUT PATH_CHECKPOINT EPOCHS PREDICT_COUNT NO_REMOVE_ISOLATED_PIXELS LOSS; echo ">>> Installing requirements"; conda run -n py38 pip install -q -r requirements.txt; diff --git a/aimodel/src/deeplabv3_plus_test_rainfall.py b/aimodel/src/deeplabv3_plus_test_rainfall.py index ca1a410..1cfa7a0 100755 --- a/aimodel/src/deeplabv3_plus_test_rainfall.py +++ b/aimodel/src/deeplabv3_plus_test_rainfall.py @@ -17,6 +17,7 @@ import matplotlib.pyplot as plt import tensorflow as tf from lib.dataset.dataset_mono import dataset_mono +from lib.ai.components.LossCrossEntropyDice import LossCrossEntropyDice IMAGE_SIZE = int(os.environ["IMAGE_SIZE"]) if "IMAGE_SIZE" in os.environ else 128 # was 512; 128 is the highest power of 2 that fits the data BATCH_SIZE = int(os.environ["BATCH_SIZE"]) if "BATCH_SIZE" in os.environ else 64 @@ -27,12 +28,12 @@ PATH_COLOURMAP = os.environ["PATH_COLOURMAP"] STEPS_PER_EPOCH = int(os.environ["STEPS_PER_EPOCH"]) if "STEPS_PER_EPOCH" in os.environ else None REMOVE_ISOLATED_PIXELS = FALSE if "NO_REMOVE_ISOLATED_PIXELS" in os.environ else True EPOCHS = int(os.environ["EPOCHS"]) if "EPOCHS" in os.environ else 25 -PREDICT_COUNT = int(os.environ["PREDICT_COUNT"]) if "PREDICT_COUNT" in os.environ else 4 - +LOSS = os.environ["LOSS"] if "LOSS" in os.environ else "cross-entropy" DIR_OUTPUT=os.environ["DIR_OUTPUT"] if "DIR_OUTPUT" in os.environ else f"output/{datetime.utcnow().date().isoformat()}_deeplabv3plus_rainfall_TEST" PATH_CHECKPOINT = os.environ["PATH_CHECKPOINT"] if "PATH_CHECKPOINT" in os.environ else None +PREDICT_COUNT = int(os.environ["PREDICT_COUNT"]) if "PREDICT_COUNT" in os.environ else 4 if not os.path.exists(DIR_OUTPUT): os.makedirs(os.path.join(DIR_OUTPUT, "checkpoints")) @@ -45,7 +46,10 @@ logger.info(f"> PATH_COLOURMAP {PATH_COLOURMAP}") logger.info(f"> STEPS_PER_EPOCH {STEPS_PER_EPOCH}") logger.info(f"> REMOVE_ISOLATED_PIXELS {REMOVE_ISOLATED_PIXELS} [NO_REMOVE_ISOLATED_PIXELS]") logger.info(f"> EPOCHS {EPOCHS}") +logger.info(f"> LOSS {LOSS}") + logger.info(f"> DIR_OUTPUT {DIR_OUTPUT}") + logger.info(f"> PATH_CHECKPOINT {PATH_CHECKPOINT}") logger.info(f"> PREDICT_COUNT {PREDICT_COUNT}") @@ -152,10 +156,17 @@ else: # ██ ██ ██ ██ ██ ██ ██ ████ ██ ██ ████ ██████ if PATH_CHECKPOINT is None: - loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) + loss_fn = None + if LOSS == "cross-entropy-dice": + loss_fn = LossCrossEntropyDice() + elif LOSS == "cross-entropy": + tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) + else: + raise Exception(f"Error: Unknown loss function '{LOSS}' (possible values: cross-entropy, cross-entropy-dice).") + model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), - loss=loss, + loss=loss_fn, metrics=["accuracy"], ) logger.info(">>> Beginning training") diff --git a/aimodel/src/lib/ai/components/LossCrossEntropyDice.py b/aimodel/src/lib/ai/components/LossCrossEntropyDice.py new file mode 100644 index 0000000..7e7713a --- /dev/null +++ b/aimodel/src/lib/ai/components/LossCrossEntropyDice.py @@ -0,0 +1,41 @@ +import math + +import tensorflow as tf + + +def dice_loss(y_true, y_pred): + """Compute Dice loss. + @source https://lars76.github.io/2018/09/27/loss-functions-for-segmentation.html#9 + Args: + y_true (tf.Tensor): The ground truth label. + y_pred (tf.Tensor): The output predicted by the model. + + Returns: + tf.Tensor: The computed Dice loss. + """ + y_pred = tf.math.sigmoid(y_pred) + numerator = 2 * tf.reduce_sum(y_true * y_pred) + denominator = tf.reduce_sum(y_true + y_pred) + + return 1 - numerator / denominator + +class LossCrossEntropyDice(tf.keras.losses.Loss): + """Cross-entropy loss and dice loss combined together into one nice neat package. + Combines the two with mean. + @source https://lars76.github.io/2018/09/27/loss-functions-for-segmentation.html#9 + """ + + def __init__(self, **kwargs): + super(LossDice, self).__init__(**kwargs) + + def call(self, y_true, y_pred): + y_true = tf.cast(y_true, tf.float32) + o = tf.nn.sigmoid_cross_entropy_with_logits(y_true, y_pred) + dice_loss(y_true, y_pred) + return tf.reduce_mean(o) + + def get_config(self): + config = super(LossDice, self).get_config() + config.update({ + + }) + return config