Add sampling

Author: Warvito

README.md

@@ -19,3 +19,12 @@ TODO LIST:
 - [ ] Add synthetic sentences based on other source of information
 - [ ] Maybe use LLM to augment the reports
 - [ ] Add warmup time for the diffusion model
+
+
+## C1
+### Uploading dataset
+To create dataset for C1, run the following command:
+
+```
+ngc dataset upload -y --desc "MIMIC dataset with dimension 512x512." --source /nfs/home/wds20/datasets/MIMIC-CXR-JPG_v2.0.0 --threads 12 scotheart
+```

src/python/testing/generate_sample_local.py

@@ -0,0 +1,105 @@
+import matplotlib.pyplot as plt
+import mlflow.pytorch
+import torch
+from generative.networks.nets import AutoencoderKL, DiffusionModelUNet
+from generative.networks.schedulers import DDIMScheduler
+from monai.config import print_config
+from monai.utils import set_determinism
+from tqdm import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer
+
+seed = 42
+set_determinism(seed=seed)
+print_config()
+
+# output_dir = Path("/media/walter/Storage/Projects/generative_cardiac/outputs/figures/same_seed")
+# output_dir.mkdir(exist_ok=True, parents=True)
+#
+stage1_old = mlflow.pytorch.load_model(
+    "/media/walter/Storage/Projects/generative_mimic/mlruns/398344666374521908/6f280de5aa634aab96e6c31eed22a62b/artifacts/final_model"
+)
+stage1 = AutoencoderKL(
+    spatial_dims=2,
+    in_channels=1,
+    out_channels=1,
+    num_channels=[64, 128, 128, 128],
+    latent_channels=3,
+    num_res_blocks=2,
+    attention_levels=[False, False, False, False],
+    with_encoder_nonlocal_attn=True,
+    with_decoder_nonlocal_attn=True,
+)
+stage1.load_state_dict(stage1_old.state_dict())
+stage1.eval()
+del stage1_old
+
+diffusion_old = mlflow.pytorch.load_model(
+    "/media/walter/Storage/Projects/generative_mimic/mlruns/411881789846457862/6f1d5a773cf5421aadd7ff787bfe7643/artifacts/final_model"
+)
+diffusion = DiffusionModelUNet(
+    spatial_dims=2,
+    in_channels=3,
+    out_channels=3,
+    num_res_blocks=2,
+    num_channels=[256, 512, 768],
+    attention_levels=[False, True, True],
+    with_conditioning=True,
+    cross_attention_dim=1024,
+    num_head_channels=[0, 512, 768],
+)
+diffusion.load_state_dict(diffusion_old.state_dict())
+diffusion.eval()
+del diffusion_old
+
+
+device = torch.device("cuda")
+diffusion = diffusion.to(device)
+stage1 = stage1.to(device)
+
+scheduler = DDIMScheduler(
+    num_train_timesteps=1000,
+    beta_start=0.0015,
+    beta_end=0.0205,
+    beta_schedule="scaled_linear",
+    prediction_type="v_prediction",
+    clip_sample=False,
+)
+scheduler.set_timesteps(200)
+
+text_encoder = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="text_encoder")
+tokenizer = CLIPTokenizer.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="tokenizer")
+
+prompt = ["", "small right-sided pleural effusion"]
+text_inputs = tokenizer(
+    prompt,
+    padding="max_length",
+    max_length=tokenizer.model_max_length,
+    truncation=True,
+    return_tensors="pt",
+)
+text_input_ids = text_inputs.input_ids
+
+prompt_embeds = text_encoder(text_input_ids.squeeze(1))
+prompt_embeds = prompt_embeds[0].to(device)
+
+guidance_scale = 7.0
+noise = torch.randn((1, 3, 64, 64)).to(device)
+
+with torch.no_grad():
+    progress_bar = tqdm(scheduler.timesteps)
+    for t in progress_bar:
+        noise_input = torch.cat([noise] * 2)
+        model_output = diffusion(
+            noise_input, timesteps=torch.Tensor((t,)).to(noise.device).long(), context=prompt_embeds
+        )
+        noise_pred_uncond, noise_pred_text = model_output.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        noise, _ = scheduler.step(noise_pred, t, noise)
+
+with torch.no_grad():
+    sample = stage1.decode_stage_2_outputs(noise / 0.3)
+
+
+plt.imshow(sample.cpu()[0, 0, :, :], cmap="gray", vmin=0, vmax=1)
+plt.show()

src/python/training/training_functions.py

@@ -4,6 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from generative.losses.adversarial_loss import PatchAdversarialLoss
 from pynvml.smi import nvidia_smi
 from tensorboardX import SummaryWriter
 from torch.cuda.amp import GradScaler, autocast
@@ -143,6 +144,8 @@ def train_epoch_aekl(
     model.train()
     discriminator.train()
 
+    adv_loss = PatchAdversarialLoss(criterion="least_squares", no_activation_leastsq=True)
+
     pbar = tqdm(enumerate(loader), total=len(loader))
     for step, x in pbar:
         images = x["image"].to(device)
@@ -157,9 +160,11 @@ def train_epoch_aekl(
             kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3])
             kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
 
-            logits_fake = discriminator(reconstruction.contiguous().float())[-1]
-            real_label = torch.ones_like(logits_fake, device=logits_fake.device)
-            generator_loss = F.mse_loss(logits_fake, real_label)
+            if adv_weight > 0:
+                logits_fake = discriminator(reconstruction.contiguous().float())[-1]
+                generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False)
+            else:
+                generator_loss = torch.tensor([0.0]).to(device)
 
             loss = l1_loss + kl_weight * kl_loss + perceptual_weight * p_loss + adv_weight * generator_loss
 
@@ -184,25 +189,26 @@ def train_epoch_aekl(
         scaler_g.update()
 
         # DISCRIMINATOR
-        optimizer_d.zero_grad(set_to_none=True)
-
-        with autocast(enabled=True):
-            logits_fake = discriminator(reconstruction.contiguous().detach())[-1]
-            fake_label = torch.zeros_like(logits_fake, device=logits_fake.device)
-            loss_d_fake = F.mse_loss(logits_fake, fake_label)
-            logits_real = discriminator(images.contiguous().detach())[-1]
-            real_label = torch.ones_like(logits_real, device=logits_real.device)
-            loss_d_real = F.mse_loss(logits_real, real_label)
-            discriminator_loss = (loss_d_fake + loss_d_real) * 0.5
-
-            d_loss = adv_weight * discriminator_loss
-            d_loss = d_loss.mean()
-
-        scaler_d.scale(d_loss).backward()
-        scaler_d.unscale_(optimizer_d)
-        torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1)
-        scaler_d.step(optimizer_d)
-        scaler_d.update()
+        if adv_weight > 0:
+            optimizer_d.zero_grad(set_to_none=True)
+
+            with autocast(enabled=True):
+                logits_fake = discriminator(reconstruction.contiguous().detach())[-1]
+                loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True)
+                logits_real = discriminator(images.contiguous().detach())[-1]
+                loss_d_real = adv_loss(logits_real, target_is_real=True, for_discriminator=True)
+                discriminator_loss = (loss_d_fake + loss_d_real) * 0.5
+
+                d_loss = adv_weight * discriminator_loss
+                d_loss = d_loss.mean()
+
+            scaler_d.scale(d_loss).backward()
+            scaler_d.unscale_(optimizer_d)
+            torch.nn.utils.clip_grad_norm_(discriminator.parameters(), 1)
+            scaler_d.step(optimizer_d)
+            scaler_d.update()
+        else:
+            discriminator_loss = torch.tensor([0.0]).to(device)
 
         losses["d_loss"] = discriminator_loss
 
@@ -241,6 +247,7 @@ def eval_aekl(
     model.eval()
     discriminator.eval()
 
+    adv_loss = PatchAdversarialLoss(criterion="least_squares", no_activation_leastsq=True)
     total_losses = OrderedDict()
     for x in loader:
         images = x["image"].to(device)
@@ -250,20 +257,24 @@ def eval_aekl(
             reconstruction, z_mu, z_sigma = model(x=images)
             l1_loss = F.l1_loss(reconstruction.float(), images.float())
             p_loss = perceptual_loss(reconstruction.float(), images.float())
-            kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3])
+            kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3, 4])
             kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
-            logits_fake = discriminator(reconstruction.contiguous().float())[-1]
-            real_label = torch.ones_like(logits_fake, device=logits_fake.device)
-            generator_loss = F.mse_loss(logits_fake, real_label)
+
+            if adv_weight > 0:
+                logits_fake = discriminator(reconstruction.contiguous().float())[-1]
+                generator_loss = adv_loss(logits_fake, target_is_real=True, for_discriminator=False)
+            else:
+                generator_loss = torch.tensor([0.0]).to(device)
 
             # DISCRIMINATOR
-            logits_fake = discriminator(reconstruction.contiguous().detach())[-1]
-            fake_label = torch.zeros_like(logits_fake, device=logits_fake.device)
-            loss_d_fake = F.mse_loss(logits_fake, fake_label)
-            logits_real = discriminator(images.contiguous().detach())[-1]
-            real_label = torch.ones_like(logits_real, device=logits_real.device)
-            loss_d_real = F.mse_loss(logits_real, real_label)
-            discriminator_loss = (loss_d_fake + loss_d_real) * 0.5
+            if adv_weight > 0:
+                logits_fake = discriminator(reconstruction.contiguous().detach())[-1]
+                loss_d_fake = adv_loss(logits_fake, target_is_real=False, for_discriminator=True)
+                logits_real = discriminator(images.contiguous().detach())[-1]
+                loss_d_real = adv_loss(logits_real, target_is_real=True, for_discriminator=True)
+                discriminator_loss = (loss_d_fake + loss_d_real) * 0.5
+            else:
+                discriminator_loss = torch.tensor([0.0]).to(device)
 
             loss = l1_loss + kl_weight * kl_loss + perceptual_weight * p_loss + adv_weight * generator_loss