seg_diff_pipeline.py

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from UNET_seg_diff import UNetModel
from mask_dataset import MaskDataset
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from tqdm import tqdm
import os
from diffusers import DDPMScheduler
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import random_split
from PIL import Image

# Time Embedding Module
num_timesteps = 1000
time_dim = 16
image_folder = "/scratch/public_scratch/Tejas_St/Talk2Car-RefSeg/images/imgs"

mask_folder="/scratch/public_scratch/Tejas_St/Talk2Car-RefSeg/masks/val_masks_new"

dataset = MaskDataset(image_folder, mask_folder)

# Split dataset into train and validation
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

batch_size = 8
train_dataloader = DataLoader(train_dataset, batch_size, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size, shuffle=False)

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

# Scheduler
scheduler = DDPMScheduler(
    num_train_timesteps=num_timesteps,
    beta_start=0.0001,
    beta_end=0.02,
    beta_schedule="linear",
)

# Forward Diffusion Process
def forward_diffusion(mask, t,scheduler):
    alpha_cumprod = scheduler.alphas_cumprod.to(device).unsqueeze(-1)
    alpha_t = alpha_cumprod[t].to(device)
    alpha_t=alpha_t.unsqueeze(-1).unsqueeze(-1)
    noise = (torch.randn_like(mask))

   
   
    mask_t = torch.sqrt(alpha_t) * mask + torch.sqrt(1 - alpha_t) * noise
    
  
    return mask_t,noise
# it seems adding common noise for both u,v makes sense, hence not changing

# def register_grad_hooks(model, threshold=1e-4):
#     def hook_fn(grad):
#         grad_mean = grad.abs().mean()  # Get mean absolute gradient
#         if grad_mean > threshold:
#             print(f"Gradient mean: {grad_mean.item()}")

#     def register_hooks_for_layers(layer):
#         if isinstance(layer, torch.nn.Conv2d) or isinstance(layer, torch.nn.Linear):
#             # Register hook on the output of the layer
#             layer.register_forward_hook(lambda module, input, output: output.register_hook(hook_fn))

#     # Register hooks for each layer
#     for name, module in model.named_modules():
#         register_hooks_for_layers(module)



# Model, optimizer, loss function
model = UNetModel(
    in_channels=1,            # Example: RGB input
    model_channels=256,        # Base channels for the model
    out_channels=1,           # Example: RGB output
    num_res_blocks=1,         # Number of residual blocks per downsample
    attention_resolutions=[4], # Attention resolutions
    dropout=0.1,              # Dropout rate
    channel_mult=(1, 2, 4),   # Channel multiplier
    conv_resample=True,       # Use convolution for resampling
    dims=2,                   # 2D UNet
    num_classes=None,         # Not class-conditional
    use_checkpoint=False,     # No gradient checkpointing
    num_heads=1,              # Number of attention heads
    use_scale_shift_norm=False, # Whether to use scale-shift normalization
).to(device)


optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()

# register_grad_hooks(model)


# Directories
os.makedirs("model_checkpoints", exist_ok=True)
os.makedirs("loss_logs", exist_ok=True)
loss_log_file = open("loss_logs/loss.txt", "a")  # Append mode for resuming training
val_log_file=open("loss_logs/val_loss.txt", "a")  # Append mode for resuming training



writer = SummaryWriter("logs")
num_epochs = 1000
save_interval = 1

# Resume Training: Load Checkpoint
start_epoch = 0
checkpoint_path = "model_checkpoints/latest_checkpoint.pth"
if os.path.exists(checkpoint_path):
    print("Loading checkpoint...")
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch'] + 1
    print(f"Resuming training from epoch {start_epoch}...")

# Training Loop
iter = start_epoch * len(train_dataloader)
val_iter=start_epoch*len(val_dataloader)
for epoch in range(start_epoch, num_epochs):
    model.train()
    epoch_loss = 0
    for im, mask in tqdm(train_dataloader):
        im=im.to(device)
        mask=mask.to(device)
        iter += 1
        optimizer.zero_grad()

        
        batch_size_local = mask.shape[0]
        t = torch.randint(0, num_timesteps, (batch_size_local,), dtype=torch.long).to(device)


        mask_t, noise = forward_diffusion(mask, t, scheduler)
        predicted_noise = model(mask_t, t, conditioned_image=im)
        

        

        loss = criterion(predicted_noise, noise)
        loss.backward()
     
        optimizer.step()

        epoch_loss += loss.item()
      
        writer.add_scalar("Loss/train_iter", loss.item(), iter)

    avg_epoch_loss = epoch_loss / len(train_dataloader)

    writer.add_scalar("Loss/train", avg_epoch_loss, iter)

    
    print(f"Epoch {epoch+1}, Training Loss: {avg_epoch_loss:.4f}")

    # Save loss to file
    loss_log_file.write(f"Epoch {epoch + 1}, Loss: {avg_epoch_loss:.4f}\n")
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'epoch': epoch,
    }, "model_checkpoints/latest_checkpoint.pth")

    # Validation
    model.eval()
    val_loss = 0
    
    with torch.no_grad():
        for im, mask in tqdm(val_dataloader):
            im=im.to(device)
            mask=mask.to(device)
            val_iter += 1

            
            batch_size_local = mask.shape[0]
            t = torch.randint(0, num_timesteps, (batch_size_local,), dtype=torch.long).to(device)


            mask_t, noise = forward_diffusion(mask, t, scheduler)
            predicted_noise = model(mask_t, t, conditioned_image=im)
            

            

            loss = criterion(predicted_noise, noise)
            val_loss += loss.item()

            writer.add_scalar("Loss/val_iter", loss.item(), val_iter)

            



            
    avg_val_loss = val_loss / len(val_dataloader)
    print(f"Epoch {epoch + 1}, Validation Loss: {avg_val_loss:.4f}")
    writer.add_scalar("Loss/val", avg_val_loss, iter)

    # Save Checkpoint
    

    if (epoch + 1) % save_interval == 0:
        torch.save(model.state_dict(), f"model_checkpoints/model_epoch_{epoch + 1}.pth")
        print(f"Model checkpoint saved for epoch {epoch + 1}")

# Final Save
torch.save(model.state_dict(), f"mode][l_checkpoints/model_epoch_{num_epochs}.pth")
loss_log_file.close()
writer.close()