This repository was archived by the owner on Feb 7, 2025. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 111
Hi all, I have already trained a diffusion model with class conditioning (via embedding layers) for weakly-supervised anomaly detection (classifier-free guidance) following the tutorial here . This was trained for 1000 epochs.
I want to take the trained weights of this diffusion model and embedding layers and train a ControlNet model using class and boundary mask condition. The tutorial here trains a ControlNet model with the pretrained weights from diffusion model checkpoint, but doesn't implement a class conditional ControlNet. I want my ControlNet to have both class and boundary mask conditioning. I used the following training and inference codes and trained the ControlNet for 100 epochs but I was not able to detect any anomaly (although the initial diffusion model was able to detect anomalies). So, I would like if someone could verify my code before I train the ControlNet for more epochs (since training this takes several days) on my machines:
Training code (DDP implementation):
#%%
import os
import time
import os
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
import sys
from monai.utils import set_determinism
from torch.amp import GradScaler, autocast
from tqdm import tqdm
import argparse
import pandas as pd
from generative.inferers import DiffusionInferer, ControlNetDiffusionInferer
from generative.networks.nets import DiffusionModelUNet, ControlNet
from generative.networks.schedulers import DDIMScheduler, DDPMScheduler
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
import torch.distributed as dist
from get_data import get_train_valid_datasets_for_controlnet
torch.multiprocessing.set_sharing_strategy("file_system")
from monai.data import DataLoader, CacheDataset
torch.backends.cudnn.benchmark = True
WORKING_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
sys.path.append(WORKING_DIR)
from utils.utils import pad_zeros_at_front, str2bool
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
#%%
def ddp_setup():
dist.init_process_group(backend='nccl', init_method="env://")
set_determinism(42)
def main_worker(models_dir, logs_dir, args):
# init_process_group
ddp_setup()
# get local rank on the GPU
local_rank = int(dist.get_rank())
if local_rank == 0:
print(f"The models will be saved in {models_dir}")
print(f"The training/validation logs will be saved in {logs_dir}")
dataset_train, dataset_valid = get_train_valid_datasets_for_controlnet(args.cache_rate)
dataset_train = dataset_train + dataset_valid
print(f'Total number of images: {dataset_train}')
sampler_train = DistributedSampler(dataset=dataset_train, shuffle=True)
dataloader_train = DataLoader(
dataset_train,
batch_size=args.batch_size,
pin_memory=True,
shuffle=False,
sampler=sampler_train,
num_workers=args.num_workers
)
trainlog_fpath = os.path.join(logs_dir, f'trainlog_gpu{local_rank}.csv')
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
embedding_dimension = 64
model = DiffusionModelUNet(
spatial_dims=2,
in_channels=1,
out_channels=1,
num_channels=(64, 64, 64),
attention_levels=(False, True, True),
num_res_blocks=1,
num_head_channels=16,
with_conditioning=True,
cross_attention_dim=embedding_dimension,
).to(device)
embed = torch.nn.Embedding(num_embeddings=3, embedding_dim=embedding_dimension, padding_idx=0).to(device)
scheduler = DDIMScheduler(num_train_timesteps=1000)
diffusion_model_path = '/data/blobfuse/diffusion2d_wsad/results/ddpm_ig/AttnLayers_011/models/chkpt_ep01000.pth'
state_dict = torch.load(diffusion_model_path, map_location=device)
model.load_state_dict(state_dict['model_state_dict'])
embed.load_state_dict(state_dict['embed_state_dict'])
controlnet = ControlNet(
spatial_dims=2,
in_channels=1,
num_channels=(64, 64, 64),
attention_levels=(False, True, True),
num_res_blocks=1,
num_head_channels=16,
conditioning_embedding_num_channels=(16,),
with_conditioning=True,
cross_attention_dim=embedding_dimension,
).to(device)
controlnet.load_state_dict(model.state_dict(), strict=False)
scheduler = DDIMScheduler(num_train_timesteps=1000)
for p in model.parameters():
p.requires_grad = False
optimizer = torch.optim.Adam(params=list(controlnet.parameters()) + list(embed.parameters()), lr=1e-5)
controlnet_inferer = ControlNetDiffusionInferer(scheduler)
start_epoch = 0
# model = DDP(model, device_ids=[device])
embed = DDP(embed, device_ids=[device])
controlnet = DDP(controlnet, device_ids=[device])
condition_dropout = 0.15
n_epochs = args.epochs
val_interval = args.val_interval
train_epoch_loss_list = []
scaler = GradScaler()
experiment_start_time = time.time()
for epoch in range(n_epochs):
epoch_start_time = time.time()
model.train()
epoch_loss = 0
sampler_train.set_epoch(epoch)
progress_bar = tqdm(enumerate(dataloader_train), total=len(dataloader_train), ncols=80)
progress_bar.set_description(f"Epoch {start_epoch + epoch + 1}")
for step, batch in progress_bar:
images, classes, boundarymasks = batch['PT'].to(device), batch['Label'].to(device), batch['BoundaryMask'].to(device)
classes = classes * (torch.rand_like(classes) > condition_dropout)
class_embedding = embed(classes.long().to(device)).unsqueeze(1)
optimizer.zero_grad(set_to_none=True)
timesteps = torch.randint(0, 1000, (len(images),)).to(device)
with autocast('cuda'):
noise = torch.randn_like(images).to(device)
noise_pred = controlnet_inferer(
inputs=images,
diffusion_model=model,
controlnet=controlnet,
noise=noise,
timesteps=timesteps,
condition=class_embedding,
cn_cond=boundarymasks
)
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({f"GPU[{local_rank}]: Train Loss": round(epoch_loss / (step + 1), 6)})
train_epoch_loss_list.append((epoch_loss / (step + 1)))
trainlog_df = pd.DataFrame(train_epoch_loss_list, columns=['Loss'])
trainlog_df.to_csv(trainlog_fpath, index=False)
if (epoch + 1) % val_interval == 0:
saved_dict = {
'controlnet_state_dict': controlnet.module.state_dict(),
'embed_state_dict': embed.module.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
}
checkpoint_fpath = os.path.join(models_dir, f'chkpt_ep{pad_zeros_at_front(start_epoch+epoch+1, 5)}.pth')
torch.save(saved_dict, checkpoint_fpath)
epoch_end_time = (time.time() - epoch_start_time)
print(f"[GPU:{local_rank}]: Epoch {start_epoch + epoch + 1} time: {round(epoch_end_time,2)} sec")
experiment_end_time = (time.time() - experiment_start_time)/(60)
print(f"[GPU:{local_rank}]: Total time: {round(experiment_end_time,2)} min")
print('Destroying process')
dist.destroy_process_group()
print('Destroyed process')
Inference code:
dataset = CacheDataset(datalist, transform=data_transforms, cache_rate=0.0)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4, drop_last=False, persistent_workers=True)
for d in dataloader:
images, gts, classes, boundarymask = d['PT'].to(device), d['GT'], d['Label'].to(device), d['BoundaryMask'].to(device)
PTs = images[:, 0].clone().detach().cpu().numpy()
GTs = gts[:, 0].clone().detach().cpu().numpy()
current_img = images.to(device)
# Noising
local_rank = 0
batch_size=1
scheduler.clip_sample = False
conditioning = torch.zeros(1).long().to(device)
class_embedding = embed(conditioning).unsqueeze(dim=1)
class_embedding_batch = class_embedding.repeat(batch_size, 1, 1)
progress_bar = tqdm(range(noise_level), desc=f'Noising batch [GPU{local_rank}]')
for i in progress_bar: # noising process for args.noise_level steps
t = i
# with autocast('cuda'):
with torch.no_grad():
t_batch = torch.Tensor([t]*batch_size).to(current_img.device)
model_output = model(current_img, timesteps=t_batch, context=class_embedding_batch)
current_img, _ = scheduler.reversed_step(model_output, t, current_img)
progress_bar.set_postfix({"timestep input": t})
latent_image = current_img
LTs = latent_image[:, 0].clone().detach().cpu().numpy()
conditioning = torch.cat([torch.zeros(1).long(), torch.ones(1).long()], dim=0).to(device)
class_embedding = embed(conditioning).unsqueeze(dim=1)
class_embedding_batch = torch.zeros(2*batch_size, 1, 64).to(device)
class_embedding_batch[0:batch_size] = class_embedding[0].unsqueeze(dim=1).repeat(batch_size, 1, 1)
class_embedding_batch[batch_size:2*batch_size] = class_embedding[1].unsqueeze(dim=1).repeat(batch_size, 1, 1)
# Denoising
from torch.amp import GradScaler, autocast
progress_bar = tqdm(range(noise_level), desc=f'Denoising batch [GPU{local_rank}]')
for i in progress_bar:
t = noise_level - i
current_img_double = torch.cat([current_img] * 2)
boundary_masks_double = torch.cat([boundarymask]*2)
# class_embedding_batch_double = class_embedding_batch.repeat(2,1,1)
with torch.no_grad():
with autocast('cuda'):
t_batch = torch.Tensor([t, t]*batch_size).to(current_img.device)
down_block_res_samples, mid_block_res_sample = controlnet(
x=current_img_double, timesteps=t_batch, controlnet_cond=boundary_masks_double, context = class_embedding_batch
)
model_output = model(
current_img_double,
timesteps=t_batch,
down_block_additional_residuals=down_block_res_samples,
mid_block_additional_residual=mid_block_res_sample,
)
noise_pred_uncond, noise_pred_text = model_output.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
current_img, _ = scheduler.step(noise_pred, t, current_img)
HLs = current_img[:, 0].clone().detach().cpu().numpy() # healthy counterfactual
ANs = abs(PTs - HLs) #anomaly map
Thanks in advance and Happy holidays!
All reactions
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment