Getting nan as generated image by using cosine schedule for DDPMScheduler

[Aman0phy]

MONAI version: 1.3.0
I want to use the cosine schedule for the DDPM scheduler
I am trying
scheduler = DDPMScheduler(num_train_timesteps=total_timesteps,schedule="cosine")

But once I switch to the cosine scheduler, during inference using the following line of code:
image = inferer.sample(input_noise=noise, diffusion_model=model, scheduler=scheduler)
The image contains the nan values for all pixels.
Can you please tell me what is the problem here?

Below are the details of the model and training loop:

Model Definition:

model = DiffusionModelUNet(
    spatial_dims=2,
    in_channels=1,
    out_channels=1,
    num_channels=(64,128, 256,512,1024),
    attention_levels=(False,False, True, True,True),
    num_res_blocks=3,
    num_head_channels=256,
)
model.to(device)

scheduler = DDPMScheduler(num_train_timesteps=total_timesteps,schedule="cosine")

optimizer = torch.optim.Adam(params=model.parameters(), lr=2.5e-5)

inferer = DiffusionInferer(scheduler)

Training loop:

if use_pretrained_model==False:
    epoch_loss_list = []
    val_epoch_loss_list = []
    memory_list = []

    scaler = GradScaler()
    total_start = time.time()
    for epoch in range(n_epochs):
        #print(f"memory usage before epoch {epoch} is {torch.cuda.mem_get_info(device)}")
        mem = torch.cuda.mem_get_info(device)
        frac = (1-(mem[0]/1000000000)/(mem[1]/1000000000))*100
        print(f"fraction gpu memory usage before epoch {epoch} is {frac:.2f}%")
        memory_list.append(frac)
        model.train()
        epoch_loss = 0
        progress_bar = tqdm(enumerate(train_loader), total=len(train_loader), ncols=70)
        progress_bar.set_description(f"Epoch {epoch}")
        for step, batch in progress_bar:
            batch_ct, batch_organmap = batch
            images = batch_ct.to(device)   
            optimizer.zero_grad(set_to_none=True)

            with autocast(enabled=True):
                # Generate random noise
                noise = torch.randn_like(images).to(device)

                # Create timesteps
                timesteps = torch.randint(
                    0, inferer.scheduler.num_train_timesteps, (images.shape[0],), device=images.device
                    ).long()

                # Get model prediction
                noise_pred = inferer(inputs=images, diffusion_model=model, noise=noise, timesteps=timesteps)

                loss = F.mse_loss(noise_pred.float(), noise.float())

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            epoch_loss += loss.item()

            progress_bar.set_postfix({"loss": epoch_loss / (step + 1)})
        epoch_loss_list.append(epoch_loss / (step + 1))

        if (epoch + 1) % val_interval == 0:
            model.eval()
            val_epoch_loss = 0
            for step, batch in enumerate(test_loader):
                batch_ct, batch_organmap = batch
                images = batch_ct.to(device)   
                with torch.no_grad():
                    with autocast(enabled=True):
                        noise = torch.randn_like(images).to(device)
                        timesteps = torch.randint(
                            0, inferer.scheduler.num_train_timesteps, (images.shape[0],), device=images.device
                        ).long()
                        noise_pred = inferer(inputs=images, diffusion_model=model, noise=noise, timesteps=timesteps)
                        val_loss = F.mse_loss(noise_pred.float(), noise.float())

                val_epoch_loss += val_loss.item()
                progress_bar.set_postfix({"val_loss": val_epoch_loss / (step + 1)})
            val_epoch_loss_list.append(val_epoch_loss / (step + 1))

            # Sampling image during training
            noise = torch.randn((1, 1, 256, 256))
            noise = noise.to(device)
            scheduler.set_timesteps(num_inference_steps=1000)
            with autocast(enabled=True):
                image = inferer.sample(input_noise=noise, diffusion_model=model, scheduler=scheduler)
            print(image.shape)
            print(image)
            plt.figure(figsize=(2, 2))
            plt.imshow(image[0, 0].cpu(), cmap="gray")
            plt.tight_layout()
            plt.axis("off")
            plt.show()

    total_time = time.time() - total_start
    print(f"train completed, total time: {total_time}.")

[marksgraham]

and can you confirm if you use linear_beta you don't get nans?

[Aman0phy]

yes, everything works perfectly with linear_beta. I also ran with sigmoid_beta, and scaled_linear_beta for testing no nans there too.
only problem is with cosine schedule

[Aman0phy]

Hi @marksgraham , Any thoughts on this ?

Thanks

[marksgraham]

I've noticed it trains fine and the samples are ok for 95% of the sampling loop, it is just the last few timesteps that cause nans to appear in the sample. I'm not sure why this is happening still, hopefully I'll have some time next week to look into this in more detail

[marksgraham]

This might be the reason, if you want to try out the code snippet i'd be interested to hear it solves your NaN issues:
#397 (comment)

[sRassmann]

I remember I had the same NaN issues and the snippet solved it for me (I first thought it was related to using FP16/mixed precision and wanted to experiment with even steeper early SNR growth, so I just implemented it like this). So the snipped should fix it.