From 09f81b074688e04170f30ce18374939e251281aa Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Mon, 28 Nov 2022 16:46:17 +0000
Subject: [PATCH] train_mono: debug this commit will generate a large amount of
 debug output.

---
 .../src/lib/ai/components/LossContrastive.py  |  7 ++++
 .../src/lib/ai/components/LossCrossentropy.py | 38 +++++++++++++++++++
 .../src/lib/ai/model_rainfallwater_mono.py    |  6 ++-
 3 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 aimodel/src/lib/ai/components/LossCrossentropy.py

diff --git a/aimodel/src/lib/ai/components/LossContrastive.py b/aimodel/src/lib/ai/components/LossContrastive.py
index 6f3b3bc..7a641a2 100644
--- a/aimodel/src/lib/ai/components/LossContrastive.py
+++ b/aimodel/src/lib/ai/components/LossContrastive.py
@@ -2,7 +2,14 @@ import math
 
 import tensorflow as tf
 
+
 class LossContrastive(tf.keras.losses.Loss):
+	"""Implements a contrastive loss function.
+	@warning: This does not function as it should.
+	Args:
+		weight_temperature (integer): The temperature weight (e.g. from LayerCheeseMultipleOut).
+		batch_size (integer): The batch size.
+	"""
 	def __init__(self, weight_temperature, batch_size):
 		super(LossContrastive, self).__init__()
 		self.batch_size = batch_size
diff --git a/aimodel/src/lib/ai/components/LossCrossentropy.py b/aimodel/src/lib/ai/components/LossCrossentropy.py
new file mode 100644
index 0000000..ef75b15
--- /dev/null
+++ b/aimodel/src/lib/ai/components/LossCrossentropy.py
@@ -0,0 +1,38 @@
+import math
+
+import tensorflow as tf
+
+
+class LossCrossentropy(tf.keras.losses.Loss):
+	"""Wraps the cross-entropy loss function because it's buggy.
+	@warning: tf.keras.losses.CategoricalCrossentropy() isn't functioning as intended during training...
+	Args:
+		batch_size (integer): The batch size (currently unused).
+	"""
+	def __init__(self, batch_size):
+		super(LossCrossentropy, self).__init__()
+		
+		self.param_batch_size = batch_size
+	
+	def call(self, y_true, y_pred):
+		result = tf.keras.metrics.categorical_crossentropy(y_true, y_pred)
+		result_reduce = tf.math.reduce_sum(result)
+		tf.print("DEBUG:TFPRINT:loss BEFORE_REDUCE", result, "AFTER_REDUCE", result_reduce)
+		return result
+	
+	def get_config(self):
+		config = super(LossCrossentropy, self).get_config()
+		config.update({
+			"batch_size": self.param_batch_size,
+		})
+		return config
+
+
+if __name__ == "__main__":
+	weight_temperature = tf.Variable(name="loss_temperature", shape=1, initial_value=tf.constant([
+		math.log(1 / 0.07)
+	]))
+	loss = LossCrossentropy(weight_temperature=weight_temperature, batch_size=64)
+	
+	tensor_input = tf.random.uniform([64, 2, 512])
+	print(loss(tf.constant(1), tensor_input))
\ No newline at end of file
diff --git a/aimodel/src/lib/ai/model_rainfallwater_mono.py b/aimodel/src/lib/ai/model_rainfallwater_mono.py
index 9b15228..4c019f4 100644
--- a/aimodel/src/lib/ai/model_rainfallwater_mono.py
+++ b/aimodel/src/lib/ai/model_rainfallwater_mono.py
@@ -6,6 +6,7 @@ import tensorflow as tf
 from .components.convnext import make_convnext
 from .components.convnext_inverse import do_convnext_inverse
 from .components.LayerStack2Image import LayerStack2Image
+from .components.LossCrossentropy import LossCrossentropy
 
 def model_rainfallwater_mono(metadata, shape_water_out, model_arch_enc="convnext_xtiny", model_arch_dec="convnext_i_xtiny", feature_dim=512, batch_size=64, water_bins=2):
 	"""Makes a new rainfall / waterdepth mono model.
@@ -71,8 +72,9 @@ def model_rainfallwater_mono(metadata, shape_water_out, model_arch_enc="convnext
 	
 	model.compile(
 		optimizer="Adam",
-		loss=tf.keras.losses.CategoricalCrossentropy(),
+		loss=LossCrossentropy(batch_size=batch_size),
+		# loss=tf.keras.losses.CategoricalCrossentropy(),
 		metrics=[tf.keras.metrics.CategoricalAccuracy()]
 	)
 	
-	return model
\ No newline at end of file
+	return model