Diffusion clean up (#605)

* Diffusion model conditioned on autoencoders

---------

Co-authored-by: Samuel Friedman, Rakesh Rathod <[email protected]>

Author: lucidtronix

ml4h/TensorMap.py

@@ -181,6 +181,8 @@ def __init__(
         # Infer loss from interpretation
         if self.loss is None and self.is_categorical() and self.shape[0] == 1:
             self.loss = 'sparse_categorical_crossentropy'
+        elif self.loss is None and self.is_categorical() and self.shape[0] == 2:
+            self.loss = 'binary_crossentropy'
         elif self.loss is None and self.is_categorical():
             self.loss = 'categorical_crossentropy'
         elif self.loss is None and self.is_continuous() and self.sentinel is not None:

ml4h/arguments.py

@@ -252,6 +252,15 @@ def parse_args():
         '--supervision_scalar', default=0.01, type=float,
         help='For `train_diffusion_supervise` mode, this weights the supervision loss from phenotype prediction on denoised data.',
     )
+    parser.add_argument('--encoder_file', help='Diffusion model encoder path for DiffAE training.')
+    parser.add_argument('--interpolate_min', type=float, default =-2.0,
+                        help='Diffusion model synthetic interpolation minimum continuous condition')
+    parser.add_argument('--interpolate_max', type=float, default =2.01,
+                        help='Diffusion model synthetic interpolation maximum continuous condition')
+    parser.add_argument('--interpolate_step', type=float, default =1.0,
+                        help='Diffusion model synthetic interpolation step size continuous condition')
+
+
     parser.add_argument(
          '--transformer_size', default=32, type=int,
          help='Number of output neurons in Transformer encoders and decoders, '
@@ -437,7 +446,8 @@ def parse_args():
     #Parent Sort enable or disable
     parser.add_argument('--parent_sort', default=True, type=lambda x: x.lower() == 'true', help='disable or enable parent_sort on output tmaps')
     #Dictionary outputs
-    parser.add_argument('--named_outputs', default=False, type=lambda x: x.lower() == 'true', help='pass output tmaps as dictionaries if true else pass as list')
+    parser.add_argument('--named_outputs', default=True, type=lambda x: x.lower() == 'true', help='pass output tmaps as dictionaries if true else pass as list')
+
     args = parser.parse_args()
     _process_args(args)
     return args

ml4h/ml4ht_integration/tensor_map.py

@@ -88,7 +88,10 @@ def __call__(self, path: str) -> Batch:
                         dependents = {dep.name: dep for dep in tm.dependent_map}
                     else:
                         dependents = {tm.dependent_map.name: tm.dependent_map}
-                out_batch[tm.output_name()] = tm.postprocess_tensor(
+                if tm in dependents:
+                    out_batch[tm.output_name()] = dependents[tm]
+                else:
+                    out_batch[tm.output_name()] = tm.postprocess_tensor(
                     tm.tensor_from_file(tm, hd5, dependents),
                     augment=self.augment, hd5=hd5,
                 )

ml4h/models/diffusion_blocks.py

@@ -9,6 +9,7 @@
 
 import keras
 from keras import layers
+from keras.models import load_model
 from keras.saving import register_keras_serializable
 
 from ml4h.defines import IMAGE_EXT
@@ -762,16 +763,22 @@ def in_paint(self, images_original, masks, diffusion_steps=64, num_rows=3, num_c
 @register_keras_serializable()
 class DiffusionController(keras.Model):
     def __init__(
-        self, tensor_map, output_maps, batch_size, widths, block_depth, conv_x, control_size,
-        attention_start, attention_heads, attention_modulo, diffusion_loss, sigmoid_beta, condition_strategy,
-        inspect_model, supervisor = None, supervision_scalar = 0.01,
+            self, tensor_map, output_maps, batch_size, widths, block_depth, conv_x, control_size,
+            attention_start, attention_heads, attention_modulo, diffusion_loss, sigmoid_beta, condition_strategy,
+            inspect_model, supervisor=None, supervision_scalar=0.01, encoder_file=None,
     ):
         super().__init__()
 
         self.input_map = tensor_map
         self.batch_size = batch_size
         self.output_maps = output_maps
-        self.control_embed_model = get_control_embed_model(self.output_maps, control_size)
+        if encoder_file:
+            self.autoencoder_control = True
+            self.control_embed_model = load_model(encoder_file, compile=False)
+            logging.info(f'loaded encoder for DiffAE at: {encoder_file}')
+        else:
+            self.autoencoder_control = False
+            self.control_embed_model = get_control_embed_model(self.output_maps, control_size)
         self.normalizer = layers.Normalization()
         self.network = get_control_network(self.input_map.shape, widths, block_depth, conv_x, control_size,
                                            attention_start, attention_heads, attention_modulo, condition_strategy)
@@ -780,7 +787,7 @@ def __init__(
         self.beta = sigmoid_beta
         self.supervisor = supervisor
         self.supervision_scalar = supervision_scalar
-        self.inspect_model = False# inspect_model
+        self.inspect_model = False  # inspect_model
 
     def get_config(self):
         config = super().get_config().copy()
@@ -796,7 +803,7 @@ def compile(self, **kwargs):
         if self.supervisor is not None:
             self.supervised_loss_tracker = keras.metrics.Mean(name="supervised_loss")
         if self.input_map.axes() == 3 and self.inspect_model:
-            self.kid = KernelInceptionDistance(name = "kid", input_shape = self.input_map.shape, kernel_image_size=299)
+            self.kid = KernelInceptionDistance(name="kid", input_shape=self.input_map.shape, kernel_image_size=299)
             self.ms_ssim = MultiScaleSSIM()
 
     @property
@@ -895,10 +902,13 @@ def generate_from_noise(self, control_embed, num_images, diffusion_steps, initia
     def train_step(self, batch):
         # normalize images to have standard deviation of 1, like the noises
         images = batch[0][self.input_map.input_name()]
-        #self.normalizer.adapt(images)
+        # self.normalizer.adapt(images)
         images = self.normalizer(images, training=True)
 
-        control_embed = self.control_embed_model(batch[1])
+        if self.autoencoder_control:
+            control_embed = self.control_embed_model(batch[0])
+        else:
+            control_embed = self.control_embed_model(batch[1])
 
         noises = tf.random.normal(shape=(self.batch_size,) + self.input_map.shape)
 
@@ -960,10 +970,13 @@ def train_step(self, batch):
     def test_step(self, batch):
         # normalize images to have standard deviation of 1, like the noises
         images = batch[0][self.input_map.input_name()]
-        #self.normalizer.adapt(images)
+        # self.normalizer.adapt(images)
         images = self.normalizer(images, training=False)
 
-        control_embed = self.control_embed_model(batch[1])
+        if self.autoencoder_control:
+            control_embed = self.control_embed_model(batch[0])
+        else:
+            control_embed = self.control_embed_model(batch[1])
 
         noises = tf.random.normal(shape=(self.batch_size,) + self.input_map.shape)
 
@@ -999,7 +1012,7 @@ def test_step(self, batch):
             supervised_loss = loss_fn(batch[1][self.output_maps[0].output_name()], supervised_preds)
             self.supervised_loss_tracker.update_state(supervised_loss)
             # Combine losses: add noise_loss and supervised_loss
-            noise_loss += self.supervision_scalar*supervised_loss
+            noise_loss += self.supervision_scalar * supervised_loss
 
         self.image_loss_tracker.update_state(image_loss)
         self.noise_loss_tracker.update_state(noise_loss)
@@ -1011,8 +1024,8 @@ def test_step(self, batch):
         if self.input_map.axes() == 3 and self.inspect_model:
             images = self.denormalize(images)
             generated_images = self.generate(control_embed,
-                num_images=self.batch_size, diffusion_steps=20
-            )
+                                             num_images=self.batch_size, diffusion_steps=20
+                                             )
             self.kid.update_state(images, generated_images)
             self.ms_ssim.update_state(images, generated_images, 255)
 
@@ -1025,7 +1038,10 @@ def call(self, batch, training=False):
           2. You can use model((noisy_images, noise_rates)) for inference
         """
         noisy_images, noise_rates = batch[0]
-        control_embed = self.control_embed_model(batch[1])
+        if self.autoencoder_control:
+            control_embed = self.control_embed_model(noisy_images)
+        else:
+            control_embed = self.control_embed_model(batch[1])
         # re-compute signal_rates
         signal_rates = tf.sqrt(1.0 - tf.square(noise_rates))
         # this returns (pred_noises, pred_images)
@@ -1063,8 +1079,8 @@ def plot_images(self, epoch=None, logs=None, num_rows=1, num_cols=4, reseed=None
         plt.close()
 
     def plot_reconstructions(
-        self, batch, diffusion_amount=0,
-        epoch=None, logs=None, num_rows=4, num_cols=4, prefix='./figures/',
+            self, batch, diffusion_amount=0,
+            epoch=None, logs=None, num_rows=4, num_cols=4, prefix='./figures/',
     ):
         images = batch[0][self.input_map.input_name()]
         self.normalizer.adapt(images)
@@ -1075,7 +1091,10 @@ def plot_reconstructions(
         # mix the images with noises accordingly
         noisy_images = signal_rates * images + noise_rates * noises
 
-        control_embed = self.control_embed_model(batch[1])
+        if self.autoencoder_control:
+            control_embed = self.control_embed_model(batch[0])
+        else:
+            control_embed = self.control_embed_model(batch[1])
 
         # use the network to separate noisy images to their components
         pred_noises, generated_images = self.denoise(
@@ -1111,10 +1130,9 @@ def plot_reconstructions(
         plt.close()
         return generated_images
 
-
     def control_plot_images(
-        self, control_batch, epoch=None, logs=None, num_rows=2, num_cols=8, reseed=None,
-        renoise=None,
+            self, control_batch, epoch=None, logs=None, num_rows=2, num_cols=8, reseed=None,
+            renoise=None,
     ):
         control_embed = self.control_embed_model(control_batch)
         # plot random generated images for visual evaluation of generation quality
@@ -1139,6 +1157,31 @@ def control_plot_images(
 
         return generated_images
 
+    def control_plot_images_embed(
+            self, control_embed, epoch=None, logs=None, num_rows=2, num_cols=8, reseed=None,
+            renoise=None,
+    ):
+        generated_images = self.generate(
+            control_embed,
+            num_images=max(self.batch_size, num_rows * num_cols),
+            diffusion_steps=plot_diffusion_steps,
+            reseed=reseed,
+            renoise=renoise,
+        )
+
+        plt.figure(figsize=(num_cols * 2.0, num_rows * 2.0), dpi=300)
+        for row in range(num_rows):
+            for col in range(num_cols):
+                index = row * num_cols + col
+                plt.subplot(num_rows, num_cols, index + 1)
+                plt.imshow(generated_images[index], cmap='gray')
+                plt.axis("off")
+        plt.tight_layout()
+        plt.show()
+        plt.close()
+
+        return generated_images
+
     def control_plot_images_noise(self, control_batch, initial_noise, epoch=None, logs=None, num_rows=2, num_cols=8):
         control_embed = self.control_embed_model(control_batch)
         # plot random generated images for visual evaluation of generation quality
@@ -1163,8 +1206,8 @@ def control_plot_images_noise(self, control_batch, initial_noise, epoch=None, lo
         return generated_images
 
     def control_plot_ecgs(
-        self, control_batch, epoch=None, logs=None, num_rows=2, num_cols=8, reseed=None,
-        renoise=None,
+            self, control_batch, epoch=None, logs=None, num_rows=2, num_cols=8, reseed=None,
+            renoise=None,
     ):
         control_embed = self.control_embed_model(control_batch)
         # plot random generated images for visual evaluation of generation quality