cleanup merge

Author: lucidtronix

ml4h/models/diffusion_blocks.py

@@ -498,8 +498,8 @@ def train_step(self, images_original):
                 noisy_images, noise_rates, signal_rates, training=True,
             )
 
-            noise_loss = self.loss(noises, pred_noises)  # used for training
-            image_loss = self.loss(images, pred_images)  # only used as metric
+            noise_loss = tf.reduce_mean(self.loss(noises, pred_noises))  # used for training
+            image_loss = tf.reduce_mean(self.loss(images, pred_images))  # only used as metric
             if self.use_sigmoid_loss:
                 signal_rates_squared = tf.square(signal_rates)
                 noise_rates_squared = tf.square(noise_rates)

ml4h/models/train_diffusion.py

@@ -28,7 +28,7 @@ def train_diffusion_model(args):
         optimizer=tf.keras.optimizers.AdamW(
             learning_rate=args.learning_rate, weight_decay=1e-4,
         ),
-        loss=tf.keras.losses.MeanAbsoluteError if args.diffusion_loss == 'mean_absolute_error' else tf.keras.losses.MeanSquaredError,
+        loss=keras.losses.MeanAbsoluteError if args.diffusion_loss == 'mean_absolute_error' else keras.losses.MeanSquaredError,
     )
     batch = next(iter(generate_train))
     for k in batch[0]:
@@ -223,7 +223,7 @@ def train_diffusion_control_model(args, supervised=False):
             args.sigmoid_beta, args.diffusion_condition_strategy, args.inspect_model,
         )
 
-    loss = tf.keras.losses.MeanAbsoluteError if args.diffusion_loss == 'mean_absolute_error' else tf.keras.losses.MeanSquaredError
+    loss = keras.losses.MeanAbsoluteError if args.diffusion_loss == 'mean_absolute_error' else keras.losses.MeanSquaredError
     model.compile(
         optimizer=tf.keras.optimizers.AdamW(
             learning_rate=args.learning_rate, weight_decay=1e-4,
@@ -348,7 +348,7 @@ def test_diffusion_control_model(args, unconditioned=False, supervised=False):
             args.sigmoid_beta, args.diffusion_condition_strategy, args.inspect_model,
         )
 
-    loss = tf.keras.losses.MeanAbsoluteError if args.diffusion_loss == 'mean_absolute_error' else tf.keras.losses.MeanSquaredError
+    loss = keras.losses.MeanAbsoluteError if args.diffusion_loss == 'mean_absolute_error' else keras.losses.MeanSquaredError
     model.compile(optimizer=tf.keras.optimizers.AdamW(learning_rate=args.learning_rate, weight_decay=1e-4), loss=loss)
     checkpoint_path = f"{args.output_folder}{args.id}/{args.id}"
     if os.path.exists(checkpoint_path+'.index'):