From 5c63136501ed4b5501bad4afcccc5ead6398ccd8 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:06:34 +0300 Subject: [PATCH 01/21] compat: fall back to native ops when CUDA extensions cannot be built utils/op/fused_act.py and utils/op/upfirdn2d.py compiled their CUDA extensions at import time via torch.utils.cpp_extension.load(). Since `backbones` imports these transitively, the whole package was unimportable on any machine without ninja and a matching CUDA toolchain -- including CPU-only installs: RuntimeError: Ninja is required to load C++ extensions Both modules already contain pure-PyTorch fallbacks (upfirdn2d_native and the CPU branch of fused_leaky_relu), but the unconditional load() ran before they could ever be reached. Build the extensions inside a try/except, warn once on failure, and route the dispatch helpers to the native path when the extension is unavailable. On machines that can build the kernels the behaviour is unchanged. Also drops a leftover debug print of the module path. --- utils/op/fused_act.py | 29 ++++++++++++++++++++--------- utils/op/upfirdn2d.py | 30 +++++++++++++++++++++--------- 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/utils/op/fused_act.py b/utils/op/fused_act.py index 29f45ed5..a3bcc8b9 100644 --- a/utils/op/fused_act.py +++ b/utils/op/fused_act.py @@ -7,6 +7,7 @@ """ import os +import warnings import torch from torch import nn @@ -16,14 +17,24 @@ module_path = os.path.dirname(__file__) -print("module_path = {}".format(module_path)) -fused = load( - "fused", - sources=[ - os.path.join(module_path, "fused_bias_act.cpp"), - os.path.join(module_path, "fused_bias_act_kernel.cu"), - ], -) + +# Building the fused CUDA kernel requires ninja and a matching CUDA toolchain. +# Neither is available on CPU-only installs, so fall back to the native PyTorch +# implementation below instead of failing at import time. +try: + fused = load( + "fused", + sources=[ + os.path.join(module_path, "fused_bias_act.cpp"), + os.path.join(module_path, "fused_bias_act_kernel.cu"), + ], + ) +except Exception as e: # pragma: no cover - depends on the local toolchain + warnings.warn( + "Could not build the fused_bias_act CUDA extension ({}). " + "Falling back to the native PyTorch implementation.".format(e) + ) + fused = None class FusedLeakyReLUFunctionBackward(Function): @@ -93,7 +104,7 @@ def forward(self, input): def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5): - if input.device.type == "cpu": + if fused is None or input.device.type == "cpu": rest_dim = [1] * (input.ndim - bias.ndim - 1) return ( F.leaky_relu( diff --git a/utils/op/upfirdn2d.py b/utils/op/upfirdn2d.py index b56fea91..0c417c06 100644 --- a/utils/op/upfirdn2d.py +++ b/utils/op/upfirdn2d.py @@ -7,6 +7,7 @@ """ import os +import warnings import torch from torch.nn import functional as F @@ -15,13 +16,24 @@ from collections import abc module_path = os.path.dirname(__file__) -upfirdn2d_op = load( - "upfirdn2d", - sources=[ - os.path.join(module_path, "upfirdn2d.cpp"), - os.path.join(module_path, "upfirdn2d_kernel.cu"), - ], -) + +# Building the upfirdn2d CUDA kernel requires ninja and a matching CUDA +# toolchain. Neither is available on CPU-only installs, so fall back to the +# native PyTorch implementation below instead of failing at import time. +try: + upfirdn2d_op = load( + "upfirdn2d", + sources=[ + os.path.join(module_path, "upfirdn2d.cpp"), + os.path.join(module_path, "upfirdn2d_kernel.cu"), + ], + ) +except Exception as e: # pragma: no cover - depends on the local toolchain + warnings.warn( + "Could not build the upfirdn2d CUDA extension ({}). " + "Falling back to the native PyTorch implementation.".format(e) + ) + upfirdn2d_op = None class UpFirDn2dBackward(Function): @@ -151,7 +163,7 @@ def backward(ctx, grad_output): def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)): - if input.device.type == "cpu": + if upfirdn2d_op is None or input.device.type == "cpu": out = upfirdn2d_native( input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1] ) @@ -173,7 +185,7 @@ def upfirdn2d_ada(input, kernel, up=1, down=1, pad=(0, 0)): if len(pad) == 2: pad = (pad[0], pad[1], pad[0], pad[1]) - if input.device.type == "cpu": + if upfirdn2d_op is None or input.device.type == "cpu": out = upfirdn2d_native(input, kernel, *up, *down, *pad) else: From 5328fedae649f84862b1afdaa4b6d9cd3d646e54 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:17:24 +0300 Subject: [PATCH 02/21] fix: remove leftover anomaly-detection call from the training loop train_syndiff() called torch.autograd.set_detect_anomaly(True) on every iteration, right before the generator backward pass. Anomaly mode records a stack trace for every autograd node and re-runs the backward with extra bookkeeping; it is a debugging aid, not something a training run should enable unconditionally. Left as-is it silently slows down every SynDiff training run. --- train.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/train.py b/train.py index e2629dea..4edbe8af 100644 --- a/train.py +++ b/train.py @@ -577,10 +577,8 @@ def train_syndiff(rank, gpu, args): #cycle loss errG1_cycle=F.l1_loss(x1_0_predict_cycle,real_data1) errG2_cycle=F.l1_loss(x2_0_predict_cycle,real_data2) - errG_cycle = errG1_cycle + errG2_cycle + errG_cycle = errG1_cycle + errG2_cycle - torch.autograd.set_detect_anomaly(True) - errG = args.lambda_l1_loss*errG_cycle + errG_adv + errG_cycle_adv + args.lambda_l1_loss*errG_L1 errG.backward() From 35cca828f32b41aa9ed71b353dc5a6fbce042442 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:26:01 +0300 Subject: [PATCH 03/21] fix: only strip the DDP 'module.' prefix when it is actually present load_checkpoint() rewrote every key as key[7:], assuming the checkpoint was written by a DistributedDataParallel-wrapped model. For a checkpoint saved without that wrapper the slice silently chops seven characters off every real parameter name, so load_state_dict() fails with a wall of unexpected / missing keys instead of loading the weights. Strip the prefix only when the key starts with it, which handles checkpoints from both DDP and single-process runs. --- test.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/test.py b/test.py index e79bcd47..6b27ada6 100644 --- a/test.py +++ b/test.py @@ -139,9 +139,13 @@ def load_checkpoint(checkpoint_dir, netG, name_of_network, epoch,device = 'cuda: checkpoint = torch.load(checkpoint_file, map_location=device) ckpt = checkpoint - - for key in list(ckpt.keys()): - ckpt[key[7:]] = ckpt.pop(key) + + # Checkpoints written by a DistributedDataParallel run carry a 'module.' + # prefix, single-process runs do not. Blindly dropping the first 7 + # characters corrupted every key of an unprefixed checkpoint. + prefix = 'module.' + ckpt = {(k[len(prefix):] if k.startswith(prefix) else k): v + for k, v in ckpt.items()} netG.load_state_dict(ckpt) netG.eval() #%% From 1380ecc4334f5e3411bedca7feadee860fbac9c5 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:26:01 +0300 Subject: [PATCH 04/21] compat: load the resume checkpoint with weights_only=False PyTorch 2.6 flipped the default of torch.load's weights_only argument to True. content.pth stores the run's argparse.Namespace next to the tensors, so --resume now aborts on every checkpoint the repo itself wrote: _pickle.UnpicklingError: Weights only load failed. ... Unsupported global: GLOBAL argparse.Namespace Load the resume checkpoint explicitly with weights_only=False. This file is produced by the training run itself, so the relaxed unpickling is contained to data the user already owns. --- train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/train.py b/train.py index 4edbe8af..4466d798 100644 --- a/train.py +++ b/train.py @@ -326,7 +326,9 @@ def train_syndiff(rank, gpu, args): if args.resume: checkpoint_file = os.path.join(exp_path, 'content.pth') - checkpoint = torch.load(checkpoint_file, map_location=device) + # content.pth stores the argparse.Namespace alongside the tensors, so it + # cannot be read under the weights_only=True default of torch>=2.6. + checkpoint = torch.load(checkpoint_file, map_location=device, weights_only=False) init_epoch = checkpoint['epoch'] epoch = init_epoch gen_diffusive_1.load_state_dict(checkpoint['gen_diffusive_1_dict']) From 27b96e7ba94a2b9672f487b7be7f80a1898c2770 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:30:12 +0300 Subject: [PATCH 05/21] feat: allow single-process and CPU-only training runs train.py hard-wired every run to CUDA + NCCL: init_processes() called torch.cuda.set_device() and dist.init_process_group(backend='nccl') unconditionally, train_syndiff() built device as 'cuda:{gpu}', and every network was wrapped in DistributedDataParallel. A single-GPU run therefore still had to stand up a NCCL process group, and a CPU-only run was impossible -- which also made the training loop untestable without a GPU. Set up a process group only when the run actually spans processes, fall back to CPU when CUDA is unavailable, and wrap in DDP only when a process group is live. Multi-process runs are unchanged: size > 1 takes exactly the same path as before. Because a single-process run now keeps the bare module, its state_dict has no 'module.' prefix. load_model_state() normalises the prefix on resume so checkpoints written by either layout keep loading. Verified end to end on CPU with the sample data: 4 epochs of training plus validation, then a --resume round-trip off the resulting content.pth. --- train.py | 103 +++++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 74 insertions(+), 29 deletions(-) diff --git a/train.py b/train.py index 4466d798..930efd5b 100644 --- a/train.py +++ b/train.py @@ -26,10 +26,43 @@ def copy_source(file, output_dir): shutil.copyfile(file, os.path.join(output_dir, os.path.basename(file))) def broadcast_params(params): + if not dist.is_initialized(): + return for param in params: dist.broadcast(param.data, src=0) +def resolve_device(gpu): + """Pick the training device, falling back to CPU when CUDA is unavailable.""" + if torch.cuda.is_available(): + return torch.device('cuda:{}'.format(gpu)) + return torch.device('cpu') + + +def maybe_ddp(model, device_ids): + """Wrap in DistributedDataParallel only when a process group is active. + + Single-process runs keep the bare module, so the saved state_dict has no + 'module.' prefix; load_checkpoint() in test.py handles both layouts. + """ + if not dist.is_initialized(): + return model + return nn.parallel.DistributedDataParallel(model, device_ids=device_ids or None) + + +def load_model_state(model, state_dict): + """Load a state_dict saved with or without the DDP 'module.' prefix. + + Whether the prefix is present depends on how many processes the run that + wrote the checkpoint used, so normalise it against the model at hand. + """ + target = model.module if isinstance(model, nn.parallel.DistributedDataParallel) else model + prefix = 'module.' + state_dict = {(k[len(prefix):] if k.startswith(prefix) else k): v + for k, v in state_dict.items()} + target.load_state_dict(state_dict) + + #%% Diffusion coefficients def var_func_vp(t, beta_min, beta_max): log_mean_coeff = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min @@ -197,10 +230,13 @@ def train_syndiff(rank, gpu, args): #rank = args.node_rank * args.num_process_per_node + gpu torch.manual_seed(args.seed + rank) - torch.cuda.manual_seed(args.seed + rank) - torch.cuda.manual_seed_all(args.seed + rank) - device = torch.device('cuda:{}'.format(gpu)) - + if torch.cuda.is_available(): + torch.cuda.manual_seed(args.seed + rank) + torch.cuda.manual_seed_all(args.seed + rank) + device = resolve_device(gpu) + # gpu_ids/device_ids are only meaningful for CUDA runs + gpu_ids = [gpu] if device.type == 'cuda' else [] + batch_size = args.batch_size nz = args.nz #latent dimension @@ -243,8 +279,8 @@ def train_syndiff(rank, gpu, args): gen_diffusive_2 = NCSNpp(args).to(device) #networks performing translation args.num_channels=1 - gen_non_diffusive_1to2 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=[gpu]) - gen_non_diffusive_2to1 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=[gpu]) + gen_non_diffusive_1to2 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=gpu_ids) + gen_non_diffusive_2to1 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=gpu_ids) disc_diffusive_1 = Discriminator_large(nc = 2, ngf = args.ngf, t_emb_dim = args.t_emb_dim, @@ -253,8 +289,8 @@ def train_syndiff(rank, gpu, args): t_emb_dim = args.t_emb_dim, act=nn.LeakyReLU(0.2)).to(device) - disc_non_diffusive_cycle1 = backbones.generator_resnet.define_D(gpu_ids=[gpu]) - disc_non_diffusive_cycle2 = backbones.generator_resnet.define_D(gpu_ids=[gpu]) + disc_non_diffusive_cycle1 = backbones.generator_resnet.define_D(gpu_ids=gpu_ids) + disc_non_diffusive_cycle2 = backbones.generator_resnet.define_D(gpu_ids=gpu_ids) broadcast_params(gen_diffusive_1.parameters()) broadcast_params(gen_diffusive_2.parameters()) @@ -299,15 +335,15 @@ def train_syndiff(rank, gpu, args): #ddp - gen_diffusive_1 = nn.parallel.DistributedDataParallel(gen_diffusive_1, device_ids=[gpu]) - gen_diffusive_2 = nn.parallel.DistributedDataParallel(gen_diffusive_2, device_ids=[gpu]) - gen_non_diffusive_1to2 = nn.parallel.DistributedDataParallel(gen_non_diffusive_1to2, device_ids=[gpu]) - gen_non_diffusive_2to1 = nn.parallel.DistributedDataParallel(gen_non_diffusive_2to1, device_ids=[gpu]) - disc_diffusive_1 = nn.parallel.DistributedDataParallel(disc_diffusive_1, device_ids=[gpu]) - disc_diffusive_2 = nn.parallel.DistributedDataParallel(disc_diffusive_2, device_ids=[gpu]) + gen_diffusive_1 = maybe_ddp(gen_diffusive_1, gpu_ids) + gen_diffusive_2 = maybe_ddp(gen_diffusive_2, gpu_ids) + gen_non_diffusive_1to2 = maybe_ddp(gen_non_diffusive_1to2, gpu_ids) + gen_non_diffusive_2to1 = maybe_ddp(gen_non_diffusive_2to1, gpu_ids) + disc_diffusive_1 = maybe_ddp(disc_diffusive_1, gpu_ids) + disc_diffusive_2 = maybe_ddp(disc_diffusive_2, gpu_ids) - disc_non_diffusive_cycle1 = nn.parallel.DistributedDataParallel(disc_non_diffusive_cycle1, device_ids=[gpu]) - disc_non_diffusive_cycle2 = nn.parallel.DistributedDataParallel(disc_non_diffusive_cycle2, device_ids=[gpu]) + disc_non_diffusive_cycle1 = maybe_ddp(disc_non_diffusive_cycle1, gpu_ids) + disc_non_diffusive_cycle2 = maybe_ddp(disc_non_diffusive_cycle2, gpu_ids) exp = args.exp output_path = args.output_path @@ -331,10 +367,10 @@ def train_syndiff(rank, gpu, args): checkpoint = torch.load(checkpoint_file, map_location=device, weights_only=False) init_epoch = checkpoint['epoch'] epoch = init_epoch - gen_diffusive_1.load_state_dict(checkpoint['gen_diffusive_1_dict']) - gen_diffusive_2.load_state_dict(checkpoint['gen_diffusive_2_dict']) - gen_non_diffusive_1to2.load_state_dict(checkpoint['gen_non_diffusive_1to2_dict']) - gen_non_diffusive_2to1.load_state_dict(checkpoint['gen_non_diffusive_2to1_dict']) + load_model_state(gen_diffusive_1, checkpoint['gen_diffusive_1_dict']) + load_model_state(gen_diffusive_2, checkpoint['gen_diffusive_2_dict']) + load_model_state(gen_non_diffusive_1to2, checkpoint['gen_non_diffusive_1to2_dict']) + load_model_state(gen_non_diffusive_2to1, checkpoint['gen_non_diffusive_2to1_dict']) # load G optimizer_gen_diffusive_1.load_state_dict(checkpoint['optimizer_gen_diffusive_1']) @@ -346,19 +382,19 @@ def train_syndiff(rank, gpu, args): optimizer_gen_non_diffusive_2to1.load_state_dict(checkpoint['optimizer_gen_non_diffusive_2to1']) scheduler_gen_non_diffusive_2to1.load_state_dict(checkpoint['scheduler_gen_non_diffusive_2to1']) # load D - disc_diffusive_1.load_state_dict(checkpoint['disc_diffusive_1_dict']) + load_model_state(disc_diffusive_1, checkpoint['disc_diffusive_1_dict']) optimizer_disc_diffusive_1.load_state_dict(checkpoint['optimizer_disc_diffusive_1']) scheduler_disc_diffusive_1.load_state_dict(checkpoint['scheduler_disc_diffusive_1']) - disc_diffusive_2.load_state_dict(checkpoint['disc_diffusive_2_dict']) + load_model_state(disc_diffusive_2, checkpoint['disc_diffusive_2_dict']) optimizer_disc_diffusive_2.load_state_dict(checkpoint['optimizer_disc_diffusive_2']) scheduler_disc_diffusive_2.load_state_dict(checkpoint['scheduler_disc_diffusive_2']) # load D_for cycle - disc_non_diffusive_cycle1.load_state_dict(checkpoint['disc_non_diffusive_cycle1_dict']) + load_model_state(disc_non_diffusive_cycle1, checkpoint['disc_non_diffusive_cycle1_dict']) optimizer_disc_non_diffusive_cycle1.load_state_dict(checkpoint['optimizer_disc_non_diffusive_cycle1']) scheduler_disc_non_diffusive_cycle1.load_state_dict(checkpoint['scheduler_disc_non_diffusive_cycle1']) - disc_non_diffusive_cycle2.load_state_dict(checkpoint['disc_non_diffusive_cycle2_dict']) + load_model_state(disc_non_diffusive_cycle2, checkpoint['disc_non_diffusive_cycle2_dict']) optimizer_disc_non_diffusive_cycle2.load_state_dict(checkpoint['optimizer_disc_non_diffusive_cycle2']) scheduler_disc_non_diffusive_cycle2.load_state_dict(checkpoint['scheduler_disc_non_diffusive_cycle2']) global_step = checkpoint['global_step'] @@ -718,15 +754,24 @@ def init_processes(rank, size, fn, args): """ Initialize the distributed environment. """ os.environ['MASTER_ADDR'] = args.master_address os.environ['MASTER_PORT'] = args.port_num - torch.cuda.set_device(args.local_rank) gpu = args.local_rank - dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=size) + if torch.cuda.is_available(): + torch.cuda.set_device(args.local_rank) + + # A process group is only needed when the run actually spans processes. + # Setting one up unconditionally made single-GPU runs depend on NCCL and + # made CPU-only runs impossible. + if size > 1: + dist.init_process_group(backend='nccl', init_method='env://', rank=rank, world_size=size) + fn(rank, gpu, args) - dist.barrier() - cleanup() + + if dist.is_initialized(): + dist.barrier() + cleanup() def cleanup(): - dist.destroy_process_group() + dist.destroy_process_group() #%% if __name__ == '__main__': parser = argparse.ArgumentParser('syndiff parameters') From 5cf53752111b817e75d6d5c3bb12d7947b409ab5 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:32:21 +0300 Subject: [PATCH 06/21] fix: size the validation metric buffers for the final epoch val_l1_loss and val_psnr_values were allocated with args.num_epoch slots along the epoch axis, but the training loop iterates over range(init_epoch, args.num_epoch + 1). The last epoch therefore indexes one past the end and the run dies right after finishing its final epoch of training -- after the checkpoint is written, but before the metrics for that epoch are recorded: File "train.py", line 711, in train_syndiff val_l1_loss[0,epoch,iteration]=abs(fake_sample1 -real_data).mean() IndexError: index 1 is out of bounds for axis 1 with size 1 Allocate args.num_epoch + 1 slots to match the loop bound. --- train.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 930efd5b..8377edcf 100644 --- a/train.py +++ b/train.py @@ -268,8 +268,10 @@ def train_syndiff(rank, gpu, args): sampler=val_sampler, drop_last = True) - val_l1_loss=np.zeros([2,args.num_epoch,len(data_loader_val)]) - val_psnr_values=np.zeros([2,args.num_epoch,len(data_loader_val)]) + # the epoch loop below runs over range(init_epoch, args.num_epoch + 1), + # so epoch args.num_epoch needs a slot too + val_l1_loss=np.zeros([2,args.num_epoch+1,len(data_loader_val)]) + val_psnr_values=np.zeros([2,args.num_epoch+1,len(data_loader_val)]) print('train data size:'+str(len(data_loader))) print('val data size:'+str(len(data_loader_val))) to_range_0_1 = lambda x: (x + 1.) / 2. From 82e2afe8065cf40efa63c1345234a677fb791790 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:32:21 +0300 Subject: [PATCH 07/21] fix: validate the 2->1 direction with gen_diffusive_2 train_syndiff() runs two validation loops. The first feeds contrast2 as the source and reconstructs contrast1 with gen_diffusive_1. The second swaps the tuple to (y_val, x_val) so that the source is contrast1 and the target is contrast2 -- but it still sampled from gen_diffusive_1. gen_diffusive_1 is the network that produces contrast1, so the second loop asked it to synthesise the other contrast and scored the result against contrast2. Every number written to val_l1_loss[1] and val_psnr_values[1] was therefore meaningless, and gen_diffusive_2 was never validated at all. Sample from gen_diffusive_2 in the second loop. --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 8377edcf..a1a99929 100644 --- a/train.py +++ b/train.py @@ -734,9 +734,9 @@ def train_syndiff(rank, gpu, args): x1_t = torch.cat((torch.randn_like(real_data),source_data),axis=1) #diffusion steps - fake_sample1 = sample_from_model(pos_coeff, gen_diffusive_1, args.num_timesteps, x1_t, T, args) + fake_sample1 = sample_from_model(pos_coeff, gen_diffusive_2, args.num_timesteps, x1_t, T, args) + - fake_sample1 = to_range_0_1(fake_sample1) ; fake_sample1 = fake_sample1/fake_sample1.mean() real_data = to_range_0_1(real_data) ; real_data = real_data/real_data.mean() From d226a61d1578f489790be0f761b61dff15002f09 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:41:52 +0300 Subject: [PATCH 08/21] compat: fall back to CPU in test.py when CUDA is unavailable sample_and_test() called torch.cuda.set_device() unconditionally, so inference aborted on any machine without a visible GPU: RuntimeError: No CUDA GPUs are available Select the CUDA device only when one exists, mirroring what train.py now does. Runs with a GPU are unaffected. --- test.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test.py b/test.py index 6b27ada6..2aed07d1 100644 --- a/test.py +++ b/test.py @@ -151,9 +151,11 @@ def load_checkpoint(checkpoint_dir, netG, name_of_network, epoch,device = 'cuda: #%% def sample_and_test(args): torch.manual_seed(42) - # device = 'cuda:0' - torch.cuda.set_device(args.gpu_chose) - device = torch.device('cuda:{}'.format(args.gpu_chose)) + if torch.cuda.is_available(): + torch.cuda.set_device(args.gpu_chose) + device = torch.device('cuda:{}'.format(args.gpu_chose)) + else: + device = torch.device('cpu') epoch_chosen=args.which_epoch to_range_0_1 = lambda x: (x + 1.) / 2. From f1864049e2cedd388c0fb4d48a61953af00c9ce4 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:41:52 +0300 Subject: [PATCH 09/21] fix: store synthesised volumes at the cropped image size syn_im1/syn_im2 were preallocated as (256, 256, N) while every sample was run through CenterCrop((256, 152)) two lines earlier, so the very first assignment aborted: syn_im1[:,:,iteration]=np.squeeze(fake_sample1.cpu().numpy()) ValueError: could not broadcast input array from shape (256,152) into shape (256,256) The crop output size is fixed, so this fired on the first test slice for every dataset -- test.py could not write its im_syn.mat at all. Collect the slices in a list and stack them once at the end, so the stored volume always matches the cropped geometry. The crop itself undoes the padding CreateDatasetSynthesis applies; its size is now exposed as --crop_h/--crop_w, defaulting to the previous IXI/BRATS values so existing commands keep producing the same output. Verified end to end on CPU: test.py now completes and writes im_syn.mat with shape (256, 152, N). --- test.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/test.py b/test.py index 2aed07d1..b00eca3f 100644 --- a/test.py +++ b/test.py @@ -186,13 +186,18 @@ def sample_and_test(args): save_dir = exp_path + "/generated_samples/epoch_{}".format(epoch_chosen) - crop = transforms.CenterCrop((256, 152)) + # CreateDatasetSynthesis pads every volume out to 256x256; this crop undoes + # that padding. The defaults match the IXI/BRATS geometry the paper used -- + # set --crop_h/--crop_w to your own slice size for other datasets. + crop = transforms.CenterCrop((args.crop_h, args.crop_w)) if not os.path.exists(save_dir): os.makedirs(save_dir) loss1 = np.zeros((1,len(data_loader))) loss2 = np.zeros((1,len(data_loader))) - syn_im1=np.zeros((256,256,len(data_loader))) - syn_im2=np.zeros((256,256,len(data_loader))) + # collected per slice and stacked afterwards, so the stored volume always + # matches the cropped image size + syn_im1=[] + syn_im2=[] for iteration, (x , y) in enumerate(data_loader): real_data = x.to(device, non_blocking=True) @@ -210,7 +215,7 @@ def sample_and_test(args): fake_sample1 = crop(fake_sample1) real_data = crop(real_data) source_data = crop(source_data) - syn_im1[:,:,iteration]=np.squeeze(fake_sample1.cpu().numpy()) + syn_im1.append(np.squeeze(fake_sample1.cpu().numpy())) loss1[0, iteration] = psnr(fake_sample1, real_data).cpu().numpy() print(str(iteration)) @@ -236,7 +241,7 @@ def sample_and_test(args): fake_sample2 = crop(fake_sample2) real_data = crop(real_data) source_data = crop(source_data) - syn_im2[:,:,iteration]=np.squeeze(fake_sample2.cpu().numpy()) + syn_im2.append(np.squeeze(fake_sample2.cpu().numpy())) loss2[0, iteration] = psnr(fake_sample2, real_data).cpu().numpy() print(str(iteration)) @@ -249,6 +254,9 @@ def sample_and_test(args): print(np.nanmean(loss2)) np.save('{}/psnr_values2.npy'.format(save_dir), loss2) + syn_im1 = np.stack(syn_im1, axis=-1) + syn_im2 = np.stack(syn_im2, axis=-1) + f = h5py.File(save_dir + '/im_syn.mat', "w") f.create_dataset('images_'+args.contrast1+'syn', data=syn_im1) f.create_dataset('images_'+args.contrast2+'syn', data=syn_im2) @@ -338,6 +346,10 @@ def sample_and_test(args): help='contrast selection for model') parser.add_argument('--contrast2', type=str, default='T2', help='contrast selection for model') + parser.add_argument('--crop_h', type=int, default=256, + help='height the padded output is cropped back to') + parser.add_argument('--crop_w', type=int, default=152, + help='width the padded output is cropped back to') parser.add_argument('--which_epoch', type=int, default=50) parser.add_argument('--gpu_chose', type=int, default=0) From 329ee84353a94278e8ef6846b0101354795a6ef3 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:44:37 +0300 Subject: [PATCH 10/21] fix: drop the dead args.num_channels assignment train_syndiff() set args.num_channels = 1 after both NCSNpp generators were already constructed, and define_G() takes its input_nc from its own default rather than from args -- so the assignment changed no network. What it did change is the Namespace that gets serialised into content.pth, which then recorded num_channels = 1 for a run whose diffusive generators were built with 2 channels. Rebuilding a model from those saved args produces a network whose weights will not load. Remove the assignment; content.pth now records the value the run actually used. --- train.py | 1 - 1 file changed, 1 deletion(-) diff --git a/train.py b/train.py index a1a99929..225eb2c7 100644 --- a/train.py +++ b/train.py @@ -280,7 +280,6 @@ def train_syndiff(rank, gpu, args): gen_diffusive_1 = NCSNpp(args).to(device) gen_diffusive_2 = NCSNpp(args).to(device) #networks performing translation - args.num_channels=1 gen_non_diffusive_1to2 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=gpu_ids) gen_non_diffusive_2to1 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=gpu_ids) From 7d4e13732ec573cbefc189b7451b7cfce3b0d1c2 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:50:07 +0300 Subject: [PATCH 11/21] fix: pad to exactly 256 and report unusable input clearly LoadDataSet() computed a single symmetric pad as int((256 - size) / 2) and applied it to both sides. For an odd difference the truncation loses a pixel, so the function silently returns something that is not 256 wide: >>> LoadDataSet(<151-wide volume>).shape (2, 1, 256, 255) Nothing checks this, so the mismatch only surfaces much later as a shape error deep inside the network. Split the padding across both sides and give the leftover pixel to the far side, so the result is exactly 256. The same function also failed unhelpfully on ordinary mistakes: a volume wider than 256 produced a negative pad ("index can't contain negative values"), a wrong variable name surfaced as a raw h5py KeyError, and a missing file as an h5py open failure that does not mention the expected data__.mat naming. Each of these now raises a message that names the file and says what was expected. Also close the HDF5 handle (it was never closed), reject inputs that are neither 3D nor 4D instead of transposing them blindly, and check that the two contrasts hold the same number of slices before they are zipped into a TensorDataset. --- dataset.py | 66 +++++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 50 insertions(+), 16 deletions(-) diff --git a/dataset.py b/dataset.py index 46a75694..b1762029 100644 --- a/dataset.py +++ b/dataset.py @@ -1,3 +1,5 @@ +import os + import torch.utils.data import numpy as np, h5py import random @@ -7,28 +9,60 @@ def CreateDatasetSynthesis(phase, input_path, contrast1 = 'T1', contrast2 = 'T2' target_file = input_path + "/data_{}_{}.mat".format(phase, contrast1) data_fs_s1=LoadDataSet(target_file) - + target_file = input_path + "/data_{}_{}.mat".format(phase, contrast2) data_fs_s2=LoadDataSet(target_file) - dataset=torch.utils.data.TensorDataset(torch.from_numpy(data_fs_s1),torch.from_numpy(data_fs_s2)) - return dataset + if data_fs_s1.shape[0] != data_fs_s2.shape[0]: + raise ValueError( + "'{}' and '{}' hold a different number of slices ({} vs {}); the two " + "contrasts must be aligned slice by slice.".format( + contrast1, contrast2, data_fs_s1.shape[0], data_fs_s2.shape[0])) + + dataset=torch.utils.data.TensorDataset(torch.from_numpy(data_fs_s1),torch.from_numpy(data_fs_s2)) + return dataset + + +#Dataset loading from load_dir and converintg to 256x256 +def LoadDataSet(load_dir, variable = 'data_fs', padding = True, Norm = True, + target_size = 256): + if not os.path.isfile(load_dir): + raise FileNotFoundError( + "No such data file: '{}'. Files are expected to be named " + "data__.mat inside --input_path.".format(load_dir)) + with h5py.File(load_dir,'r') as f: + if variable not in f: + raise KeyError( + "'{}' does not contain a '{}' variable (found: {}).".format( + load_dir, variable, ', '.join(f.keys()) or 'nothing')) + raw = np.array(f[variable]) -#Dataset loading from load_dir and converintg to 256x256 -def LoadDataSet(load_dir, variable = 'data_fs', padding = True, Norm = True): - f = h5py.File(load_dir,'r') - if np.array(f[variable]).ndim==3: - data=np.expand_dims(np.transpose(np.array(f[variable]),(0,2,1)),axis=1) + if raw.ndim==3: + data=np.expand_dims(np.transpose(raw,(0,2,1)),axis=1) + elif raw.ndim==4: + data=np.transpose(raw,(1,0,3,2)) else: - data=np.transpose(np.array(f[variable]),(1,0,3,2)) - data=data.astype(np.float32) + raise ValueError( + "'{}' has {} dimensions; expected 3 (slices, width, height) or 4.".format( + load_dir, raw.ndim)) + data=data.astype(np.float32) if padding: - pad_x=int((256-data.shape[2])/2) - pad_y=int((256-data.shape[3])/2) - print('padding in x-y with:'+str(pad_x)+'-'+str(pad_y)) - data=np.pad(data,((0,0),(0,0),(pad_x,pad_x),(pad_y,pad_y))) - if Norm: - data=(data-0.5)/0.5 + pads = [] + for axis in (2, 3): + size = data.shape[axis] + if size > target_size: + raise ValueError( + "'{}' has a {}x{} image size, which does not fit the {}x{} grid " + "SynDiff pads to. Crop or resample the volume first.".format( + load_dir, data.shape[2], data.shape[3], target_size, target_size)) + # split the padding across both sides; the extra pixel of an odd + # difference goes to the far side so the result is exactly target_size + before = (target_size - size) // 2 + pads.append((before, target_size - size - before)) + print('padding in x-y with:'+str(pads[0])+'-'+str(pads[1])) + data=np.pad(data,((0,0),(0,0),pads[0],pads[1])) + if Norm: + data=(data-0.5)/0.5 return data From 699768577edfe4d624606169900095351278543a Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:53:50 +0300 Subject: [PATCH 12/21] perf: run the generators under no_grad during discriminator updates Both discriminator phases ran the generators with autograd enabled and fed the results straight into the discriminators. Nothing detached them, so errD_fake.backward() and errD_cycle_fake.backward() propagated all the way back through gen_diffusive_1/2 and gen_non_diffusive_1to2/2to1, filling their .grad buffers -- which the following gen_*.zero_grad() then threw away. Every iteration paid for a full backward pass through four generators whose gradients were discarded by construction. Wrap those forward passes in torch.no_grad(). The discriminator update is unchanged mathematically: it never wanted generator gradients. Measured on one discriminator branch (256x256, batch 1, GPU), averaged over 10 iterations after warm-up: no_grad=True 67.3 ms/iter 0.14 GB transient no_grad=False 156.4 ms/iter 0.56 GB transient Peak memory for the whole training step is unchanged -- that peak is set by the generator update, which legitimately keeps its graph. The win here is wasted compute in the discriminator phases, not headroom. --- train.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/train.py b/train.py index 225eb2c7..01cd12e1 100644 --- a/train.py +++ b/train.py @@ -486,14 +486,17 @@ def train_syndiff(rank, gpu, args): latent_z1 = torch.randn(batch_size, nz, device=device) latent_z2 = torch.randn(batch_size, nz, device=device) - x1_0_predict = gen_non_diffusive_2to1(real_data2) - x2_0_predict = gen_non_diffusive_1to2(real_data1) - #x_tp1 is concatenated with source contrast and x_0_predict is predicted - x1_0_predict_diff = gen_diffusive_1(torch.cat((x1_tp1.detach(),x2_0_predict),axis=1), t1, latent_z1) - x2_0_predict_diff = gen_diffusive_2(torch.cat((x2_tp1.detach(),x1_0_predict),axis=1), t2, latent_z2) - #sampling q(x_t | x_0_predict, x_t+1) - x1_pos_sample = sample_posterior(pos_coeff, x1_0_predict_diff[:,[0],:], x1_tp1, t1) - x2_pos_sample = sample_posterior(pos_coeff, x2_0_predict_diff[:,[0],:], x2_tp1, t2) + # only D is updated here, so the generators run without building a + # graph; their gradients would be discarded by zero_grad() below + with torch.no_grad(): + x1_0_predict = gen_non_diffusive_2to1(real_data2) + x2_0_predict = gen_non_diffusive_1to2(real_data1) + #x_tp1 is concatenated with source contrast and x_0_predict is predicted + x1_0_predict_diff = gen_diffusive_1(torch.cat((x1_tp1.detach(),x2_0_predict),axis=1), t1, latent_z1) + x2_0_predict_diff = gen_diffusive_2(torch.cat((x2_tp1.detach(),x1_0_predict),axis=1), t2, latent_z2) + #sampling q(x_t | x_0_predict, x_t+1) + x1_pos_sample = sample_posterior(pos_coeff, x1_0_predict_diff[:,[0],:], x1_tp1, t1) + x2_pos_sample = sample_posterior(pos_coeff, x2_0_predict_diff[:,[0],:], x2_tp1, t2) #D output for fake sample x_pos_sample output1 = disc_diffusive_1(x1_pos_sample, t1, x1_tp1.detach()).view(-1) output2 = disc_diffusive_2(x2_pos_sample, t2, x2_tp1.detach()).view(-1) @@ -528,8 +531,9 @@ def train_syndiff(rank, gpu, args): errD_cycle_real.backward(retain_graph=True) # train with fake - x1_0_predict = gen_non_diffusive_2to1(real_data2) - x2_0_predict = gen_non_diffusive_1to2(real_data1) + with torch.no_grad(): + x1_0_predict = gen_non_diffusive_2to1(real_data2) + x2_0_predict = gen_non_diffusive_1to2(real_data1) D_cycle1_fake = disc_non_diffusive_cycle1(x1_0_predict).view(-1) D_cycle2_fake = disc_non_diffusive_cycle2(x2_0_predict).view(-1) From 020ff2f5e656d4a5f8466fe7c860de9925d27c91 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 12:57:33 +0300 Subject: [PATCH 13/21] fix: stop wrapping the translation networks in DataParallel init_net() wrapped every network it built in torch.nn.DataParallel, and train.py then wrapped the result in DistributedDataParallel. The two do not compose: DataParallel re-scatters each batch across devices inside a module that DDP already owns. train.py only ever passes a single device, so in practice the inner wrapper did no parallel work at all -- it just added a per-call replication step and a second 'module.' level to every key of the saved state_dict. Move the network to its device and leave the parallelism to DDP. Checkpoints written before this change carry 'module.module.' on the translation networks, so both loaders now strip the prefix repeatedly rather than once, and load checkpoints from any of the three layouts. Regression-checked on CPU with the sample data: training, --resume off the resulting content.pth, and test.py inference all complete. --- backbones/generator_resnet.py | 1 - test.py | 16 ++++++++++------ train.py | 25 +++++++++++++++++-------- 3 files changed, 27 insertions(+), 15 deletions(-) diff --git a/backbones/generator_resnet.py b/backbones/generator_resnet.py index d0f2f238..13590a39 100644 --- a/backbones/generator_resnet.py +++ b/backbones/generator_resnet.py @@ -96,7 +96,6 @@ def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[]): if len(gpu_ids) > 0: assert(torch.cuda.is_available()) net.to(gpu_ids[0]) - net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs init_weights(net, init_type, init_gain=init_gain) return net diff --git a/test.py b/test.py index b00eca3f..593ed422 100644 --- a/test.py +++ b/test.py @@ -140,13 +140,17 @@ def load_checkpoint(checkpoint_dir, netG, name_of_network, epoch,device = 'cuda: checkpoint = torch.load(checkpoint_file, map_location=device) ckpt = checkpoint - # Checkpoints written by a DistributedDataParallel run carry a 'module.' - # prefix, single-process runs do not. Blindly dropping the first 7 - # characters corrupted every key of an unprefixed checkpoint. + # Checkpoints carry one 'module.' prefix per parallel wrapper the saving + # run used: none for a single-process run, one for DistributedDataParallel. + # Blindly dropping the first 7 characters corrupted every key of an + # unprefixed checkpoint. prefix = 'module.' - ckpt = {(k[len(prefix):] if k.startswith(prefix) else k): v - for k, v in ckpt.items()} - netG.load_state_dict(ckpt) + normalised = {} + for key, value in ckpt.items(): + while key.startswith(prefix): + key = key[len(prefix):] + normalised[key] = value + netG.load_state_dict(normalised) netG.eval() #%% def sample_and_test(args): diff --git a/train.py b/train.py index 01cd12e1..39f7be32 100644 --- a/train.py +++ b/train.py @@ -50,17 +50,26 @@ def maybe_ddp(model, device_ids): return nn.parallel.DistributedDataParallel(model, device_ids=device_ids or None) -def load_model_state(model, state_dict): - """Load a state_dict saved with or without the DDP 'module.' prefix. +def strip_module_prefix(state_dict): + """Drop the 'module.' prefixes that parallel wrappers add to every key. - Whether the prefix is present depends on how many processes the run that - wrote the checkpoint used, so normalise it against the model at hand. + Checkpoints carry one prefix per wrapper the saving run used: none for a + single-process run, one for DistributedDataParallel, and two for the + DDP-over-DataParallel nesting older revisions produced. """ - target = model.module if isinstance(model, nn.parallel.DistributedDataParallel) else model prefix = 'module.' - state_dict = {(k[len(prefix):] if k.startswith(prefix) else k): v - for k, v in state_dict.items()} - target.load_state_dict(state_dict) + stripped = {} + for key, value in state_dict.items(): + while key.startswith(prefix): + key = key[len(prefix):] + stripped[key] = value + return stripped + + +def load_model_state(model, state_dict): + """Load a state_dict saved with or without the DDP 'module.' prefix.""" + target = model.module if isinstance(model, nn.parallel.DistributedDataParallel) else model + target.load_state_dict(strip_module_prefix(state_dict)) #%% Diffusion coefficients From f36c55d4a5c978364e452bc3de081ca4db5ca331 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:04:21 +0300 Subject: [PATCH 14/21] fix: make the minibatch-stddev grouping divide the batch Both discriminators reshape the last feature map into view(group, -1, stddev_feat, channel // stddev_feat, height, width) with group = min(batch, 4). That view is only valid when group divides the batch, so any batch size that is neither <= 4 nor a multiple of 4 aborts: batch=5 RuntimeError: shape '[4, -1, 1, 64, 2, 2]' is invalid for input of size 1280 batch=6 RuntimeError: shape '[4, -1, 1, 64, 2, 2]' is invalid for input of size 1536 batch=7 RuntimeError: shape '[4, -1, 1, 64, 2, 2]' is invalid for input of size 1792 The published command uses --batch_size 1 so this never surfaced there, but it makes several perfectly ordinary batch sizes unusable. Step the group size down to the largest value that divides the batch. Every batch size that worked before keeps its previous grouping (1, 2, 3, 4 and multiples of 4 are unchanged); the ones that crashed now run. --- backbones/discriminator.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/backbones/discriminator.py b/backbones/discriminator.py index 9a49298d..8e45eba2 100644 --- a/backbones/discriminator.py +++ b/backbones/discriminator.py @@ -149,7 +149,11 @@ def forward(self, x, t, x_t): out = self.conv4(h3,t_embed) batch, channel, height, width = out.shape + # the view below splits the batch into groups of this size, so it has to + # divide the batch evenly; fall back to the largest size that does group = min(batch, self.stddev_group) + while batch % group != 0: + group -= 1 stddev = out.view( group, -1, self.stddev_feat, channel // self.stddev_feat, height, width ) @@ -221,7 +225,11 @@ def forward(self, x, t, x_t): out = self.conv6(h,t_embed) batch, channel, height, width = out.shape + # the view below splits the batch into groups of this size, so it has to + # divide the batch evenly; fall back to the largest size that does group = min(batch, self.stddev_group) + while batch % group != 0: + group -= 1 stddev = out.view( group, -1, self.stddev_feat, channel // self.stddev_feat, height, width ) From e1cca7624b9220f79e51e2001b8a756716cf84d6 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:04:21 +0300 Subject: [PATCH 15/21] test: add a CPU test suite for the diffusion, model and data paths The repository had no tests, which is why several of the defects fixed in this branch could sit in released code: nothing ever checked the values they produced. This adds a suite that runs on CPU in a few seconds. tests/test_diffusion.py schedule and posterior coefficient properties -- the variance-preserving identity, monotone signal decay, non-negative posterior variance, the collapse to x_0 at t == 0, and that sample_posterior is deterministic there and stochastic elsewhere. tests/test_models.py shapes and gradients for NCSNpp, both discriminators and the translation networks; discriminators are exercised across batch sizes 1..8; NCSNpp must reduce a reconstruction loss over a handful of optimiser steps. tests/test_dataset.py padding reaches exactly 256 for odd and even widths, stays centred, normalises to [-1, 1], and each rejected input raises a message that names the problem. tests/test_checkpoint.py state_dicts load through all three 'module.' layouts. tests/test_ops.py the CUDA kernels agree with their native fallbacks, forward and backward; skipped automatically where the extensions cannot build. 59 tests pass with the CUDA extensions built, 54 pass with 5 skipped on a CPU-only install. --- tests/conftest.py | 23 +++++++ tests/test_checkpoint.py | 52 +++++++++++++++ tests/test_dataset.py | 88 ++++++++++++++++++++++++++ tests/test_diffusion.py | 113 +++++++++++++++++++++++++++++++++ tests/test_models.py | 132 +++++++++++++++++++++++++++++++++++++++ tests/test_ops.py | 72 +++++++++++++++++++++ 6 files changed, 480 insertions(+) create mode 100644 tests/conftest.py create mode 100644 tests/test_checkpoint.py create mode 100644 tests/test_dataset.py create mode 100644 tests/test_diffusion.py create mode 100644 tests/test_models.py create mode 100644 tests/test_ops.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..0320900d --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +import argparse +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +@pytest.fixture(scope="session") +def small_config(): + """A tiny NCSNpp/diffusion config so the tests stay fast on CPU.""" + return argparse.Namespace( + image_size=32, num_channels=2, centered=True, num_channels_dae=8, + n_mlp=2, ch_mult=[1, 2], num_res_blocks=1, attn_resolutions=(16,), + dropout=0., resamp_with_conv=True, conditional=True, fir=True, + fir_kernel=[1, 3, 3, 1], skip_rescale=True, resblock_type='biggan', + progressive='none', progressive_input='residual', + progressive_combine='sum', embedding_type='positional', + fourier_scale=16., not_use_tanh=False, nz=8, z_emb_dim=16, + t_emb_dim=16, ngf=8, num_timesteps=4, beta_min=0.1, beta_max=20., + use_geometric=False, + ) diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py new file mode 100644 index 00000000..c5b9ec8f --- /dev/null +++ b/tests/test_checkpoint.py @@ -0,0 +1,52 @@ +"""Checkpoints must survive the round trip between parallel wrappers.""" +import torch +import torch.nn as nn + +from train import load_model_state, strip_module_prefix + + +class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(1, 2, 3, padding=1) + + +def test_strips_no_prefix_from_a_single_process_checkpoint(): + assert list(strip_module_prefix({'conv.weight': 1})) == ['conv.weight'] + + +def test_strips_one_prefix_from_a_ddp_checkpoint(): + assert list(strip_module_prefix({'module.conv.weight': 1})) == ['conv.weight'] + + +def test_strips_both_prefixes_from_a_ddp_over_dataparallel_checkpoint(): + """Revisions that wrapped DataParallel inside DDP wrote two levels.""" + assert list(strip_module_prefix({'module.module.conv.weight': 1})) == ['conv.weight'] + + +def test_prefix_stripping_preserves_values_and_arity(): + state = {'module.a': torch.zeros(2), 'module.b': torch.ones(3)} + out = strip_module_prefix(state) + assert set(out) == {'a', 'b'} + assert torch.equal(out['b'], torch.ones(3)) + + +def test_load_model_state_accepts_every_layout(): + model = Tiny() + reference = Tiny() + for prefix in ('', 'module.', 'module.module.'): + state = {prefix + k: v for k, v in reference.state_dict().items()} + load_model_state(model, state) + for key, value in reference.state_dict().items(): + assert torch.equal(model.state_dict()[key], value) + + +def test_load_model_state_round_trips_through_a_real_save(tmp_path): + saved = Tiny() + path = tmp_path / 'ckpt.pth' + torch.save(saved.state_dict(), path) + + loaded = Tiny() + load_model_state(loaded, torch.load(path, map_location='cpu')) + for key, value in saved.state_dict().items(): + assert torch.equal(loaded.state_dict()[key], value) diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 00000000..d1f2e366 --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,88 @@ +"""Loading, padding and normalisation of the .mat volumes.""" +import h5py +import numpy as np +import pytest +import torch + +from dataset import CreateDatasetSynthesis, LoadDataSet + + +def write_mat(path, array, variable='data_fs'): + with h5py.File(path, 'w') as f: + f.create_dataset(variable, data=array) + return str(path) + + +def test_loads_and_normalises_to_minus_one_one(tmp_path): + raw = np.linspace(0, 1, 2 * 152 * 256, dtype=np.float32).reshape(2, 152, 256) + data = LoadDataSet(write_mat(tmp_path / 'a.mat', raw)) + assert data.shape == (2, 1, 256, 256) + assert data.dtype == np.float32 + # (x - 0.5) / 0.5 maps [0, 1] onto [-1, 1] + assert data.min() >= -1.0 - 1e-6 and data.max() <= 1.0 + 1e-6 + + +def test_padding_can_be_disabled(tmp_path): + raw = np.zeros((2, 152, 256), dtype=np.float32) + data = LoadDataSet(write_mat(tmp_path / 'a.mat', raw), padding=False) + assert data.shape == (2, 1, 256, 152) + + +@pytest.mark.parametrize("width", [151, 152, 153, 255, 256]) +def test_padding_always_reaches_the_target_size(tmp_path, width): + """An odd width used to lose a pixel to integer truncation.""" + raw = np.zeros((2, width, 256), dtype=np.float32) + data = LoadDataSet(write_mat(tmp_path / 'a.mat', raw)) + assert data.shape[2:] == (256, 256) + + +def test_padding_keeps_the_image_centred(tmp_path): + raw = np.ones((1, 152, 256), dtype=np.float32) + data = LoadDataSet(write_mat(tmp_path / 'a.mat', raw), Norm=False) + column_sums = data[0, 0].sum(axis=0) + filled = np.nonzero(column_sums)[0] + assert filled.size == 152 + # at most one pixel of asymmetry, and it goes to the far side + assert abs(filled[0] - (256 - 1 - filled[-1])) <= 1 + + +def test_oversized_input_is_rejected_with_a_useful_message(tmp_path): + raw = np.zeros((2, 300, 256), dtype=np.float32) + with pytest.raises(ValueError, match="does not fit"): + LoadDataSet(write_mat(tmp_path / 'a.mat', raw)) + + +def test_missing_file_names_the_expected_layout(tmp_path): + with pytest.raises(FileNotFoundError, match="data__"): + LoadDataSet(str(tmp_path / 'nope.mat')) + + +def test_wrong_variable_lists_what_the_file_holds(tmp_path): + path = write_mat(tmp_path / 'a.mat', np.zeros((2, 152, 256), np.float32), + variable='images') + with pytest.raises(KeyError, match="images"): + LoadDataSet(path) + + +def test_unexpected_dimensionality_is_rejected(tmp_path): + path = write_mat(tmp_path / 'a.mat', np.zeros((2, 152), np.float32)) + with pytest.raises(ValueError, match="dimensions"): + LoadDataSet(path) + + +def test_create_dataset_pairs_the_two_contrasts(tmp_path): + for contrast in ('T1', 'T2'): + write_mat(tmp_path / 'data_train_{}.mat'.format(contrast), + np.zeros((3, 152, 256), np.float32)) + dataset = CreateDatasetSynthesis('train', str(tmp_path), 'T1', 'T2') + assert len(dataset) == 3 + first, second = dataset[0] + assert isinstance(first, torch.Tensor) and first.shape == (1, 256, 256) + assert second.shape == (1, 256, 256) + + +def test_mismatched_slice_counts_are_rejected(tmp_path): + write_mat(tmp_path / 'data_train_T1.mat', np.zeros((4, 152, 256), np.float32)) + write_mat(tmp_path / 'data_train_T2.mat', np.zeros((3, 152, 256), np.float32)) + with pytest.raises(ValueError, match="different number of slices"): + CreateDatasetSynthesis('train', str(tmp_path), 'T1', 'T2') diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py new file mode 100644 index 00000000..f5e6b638 --- /dev/null +++ b/tests/test_diffusion.py @@ -0,0 +1,113 @@ +"""Properties the forward/reverse diffusion coefficients have to satisfy.""" +import numpy as np +import pytest +import torch + +from train import (Diffusion_Coefficients, Posterior_Coefficients, extract, + get_sigma_schedule, get_time_schedule, q_sample_pairs, + sample_posterior) + + +def test_sigma_schedule_shapes_and_bounds(small_config): + sigmas, a_s, betas = get_sigma_schedule(small_config, device='cpu') + n = small_config.num_timesteps + assert sigmas.shape == a_s.shape == betas.shape == (n + 1,) + assert torch.all(betas >= 0) and torch.all(betas <= 1) + # a_s = sqrt(1 - beta) is the per-step signal retention + assert torch.allclose(a_s, torch.sqrt(1 - betas), atol=1e-6) + assert torch.allclose(sigmas, betas ** 0.5, atol=1e-6) + + +def test_variance_preserving_identity(small_config): + """a_s_cum^2 + sigmas_cum^2 == 1 -- the VP property q_sample relies on.""" + coeff = Diffusion_Coefficients(small_config, device='cpu') + total = coeff.a_s_cum ** 2 + coeff.sigmas_cum ** 2 + assert torch.allclose(total, torch.ones_like(total), atol=1e-5) + + +def test_cumulative_signal_decays(small_config): + coeff = Diffusion_Coefficients(small_config, device='cpu') + a_s_cum = coeff.a_s_cum + assert torch.all(a_s_cum[1:] <= a_s_cum[:-1] + 1e-6), "signal must not grow with t" + assert torch.all(coeff.sigmas_cum[1:] >= coeff.sigmas_cum[:-1] - 1e-6) + + +def test_posterior_variance_is_positive(small_config): + pos = Posterior_Coefficients(small_config, device='cpu') + assert torch.all(pos.posterior_variance >= 0) + assert torch.all(torch.isfinite(pos.posterior_log_variance_clipped)) + assert torch.all(torch.isfinite(pos.posterior_mean_coef1)) + assert torch.all(torch.isfinite(pos.posterior_mean_coef2)) + + +def test_posterior_collapses_to_x0_at_t_zero(small_config): + """At t == 0 the posterior mean is exactly x_0: coef1 == 1, coef2 == 0.""" + pos = Posterior_Coefficients(small_config, device='cpu') + assert pos.posterior_mean_coef1[0] == pytest.approx(1.0, abs=1e-4) + assert pos.posterior_mean_coef2[0] == pytest.approx(0.0, abs=1e-4) + + +def test_posterior_mean_coefficients_are_non_negative(small_config): + pos = Posterior_Coefficients(small_config, device='cpu') + assert torch.all(pos.posterior_mean_coef1 >= 0) + assert torch.all(pos.posterior_mean_coef2 >= 0) + # weight on x_0 falls off as the step index grows + c1 = pos.posterior_mean_coef1 + assert torch.all(c1[1:] <= c1[:-1] + 1e-6) + + +def test_posterior_variance_never_exceeds_beta(small_config): + pos = Posterior_Coefficients(small_config, device='cpu') + assert torch.all(pos.posterior_variance <= pos.betas + 1e-6) + + +def test_reciprocal_alpha_helpers_are_consistent(small_config): + pos = Posterior_Coefficients(small_config, device='cpu') + product = pos.sqrt_recip_alphas_cumprod * pos.sqrt_alphas_cumprod + assert torch.allclose(product, torch.ones_like(product), atol=1e-5) + + +def test_time_schedule_is_increasing_and_bounded(small_config): + T = get_time_schedule(small_config, device='cpu') + assert T.shape == (small_config.num_timesteps + 1,) + assert torch.all(T[1:] > T[:-1]) + assert T[0] > 0 and T[-1] <= 1.0 + + +def test_extract_selects_per_sample_coefficients(): + values = torch.tensor([10., 20., 30.]) + t = torch.tensor([2, 0]) + out = extract(values, t, (2, 1, 4, 4)) + assert out.shape == (2, 1, 1, 1) + assert out.flatten().tolist() == [30., 10.] + + +def test_q_sample_pairs_keeps_shape_and_is_noisier_at_t_plus_one(small_config): + coeff = Diffusion_Coefficients(small_config, device='cpu') + torch.manual_seed(0) + x0 = torch.randn(4, 1, 8, 8) + t = torch.zeros(4, dtype=torch.int64) + x_t, x_tp1 = q_sample_pairs(coeff, x0, t) + assert x_t.shape == x_tp1.shape == x0.shape + # at t=0 the pair is one diffusion step apart, so x_tp1 is further from x0 + assert (x_tp1 - x0).abs().mean() > (x_t - x0).abs().mean() + + +def test_sample_posterior_is_deterministic_at_t_zero(small_config): + """The nonzero_mask must suppress the noise term for t == 0.""" + pos = Posterior_Coefficients(small_config, device='cpu') + x0 = torch.randn(3, 1, 8, 8) + xt = torch.randn(3, 1, 8, 8) + t = torch.zeros(3, dtype=torch.int64) + first = sample_posterior(pos, x0, xt, t) + second = sample_posterior(pos, x0, xt, t) + assert torch.allclose(first, second) + + +def test_sample_posterior_is_stochastic_for_positive_t(small_config): + pos = Posterior_Coefficients(small_config, device='cpu') + x0 = torch.randn(3, 1, 8, 8) + xt = torch.randn(3, 1, 8, 8) + t = torch.full((3,), small_config.num_timesteps - 1, dtype=torch.int64) + assert not torch.allclose(sample_posterior(pos, x0, xt, t), + sample_posterior(pos, x0, xt, t)) diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 00000000..45d99c78 --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,132 @@ +"""Shape and gradient smoke tests for the networks, small enough to run on CPU.""" +import pytest +import torch +import torch.nn as nn + +import backbones.generator_resnet as generator_resnet +from backbones.discriminator import Discriminator_large, Discriminator_small +from backbones.ncsnpp_generator_adagn import NCSNpp + + +def test_ncsnpp_roundtrip_shape_and_backward(small_config): + net = NCSNpp(small_config) + x = torch.randn(2, small_config.num_channels, + small_config.image_size, small_config.image_size) + t = torch.randint(0, small_config.num_timesteps, (2,)) + z = torch.randn(2, small_config.nz) + + out = net(x, t, z) + assert out.shape == x.shape + # the default config keeps the tanh, so outputs stay in [-1, 1] + assert out.abs().max() <= 1.0 + + out.sum().backward() + assert any(p.grad is not None and torch.isfinite(p.grad).all() + for p in net.parameters()) + + +def test_untrained_output_is_near_zero(small_config): + """NCSN++ zero-initialises its residual and output convolutions on purpose. + + Pinning this down documents why an untrained model cannot be probed for + timestep/latent sensitivity: every conditioned branch starts at zero. + """ + torch.manual_seed(0) + net = NCSNpp(small_config).eval() + x = torch.randn(1, small_config.num_channels, + small_config.image_size, small_config.image_size) + with torch.no_grad(): + out = net(x, torch.zeros(1, dtype=torch.int64), torch.randn(1, small_config.nz)) + assert out.abs().max() < 1e-3 + + +def test_timestep_embedding_separates_timesteps(small_config): + """The conditioning pathway itself: distinct t must give distinct embeddings.""" + from backbones.layers import get_timestep_embedding + + emb = get_timestep_embedding(torch.arange(small_config.num_timesteps), + small_config.num_channels_dae) + assert emb.shape == (small_config.num_timesteps, small_config.num_channels_dae) + assert torch.isfinite(emb).all() + for i in range(emb.shape[0]): + for j in range(i + 1, emb.shape[0]): + assert not torch.allclose(emb[i], emb[j]), "t={} and t={} collide".format(i, j) + + +def test_z_transform_responds_to_the_latent(small_config): + """The latent mapping network must not collapse different z to one code.""" + torch.manual_seed(0) + net = NCSNpp(small_config).eval() + with torch.no_grad(): + first = net.z_transform(torch.randn(1, small_config.nz)) + second = net.z_transform(torch.randn(1, small_config.nz)) + assert first.shape == (1, small_config.z_emb_dim) + assert not torch.allclose(first, second) + + +def test_ncsnpp_can_reduce_a_reconstruction_loss(small_config): + """A few optimiser steps must move the loss: catches a dead training path.""" + torch.manual_seed(0) + net = NCSNpp(small_config) + opt = torch.optim.Adam(net.parameters(), lr=1e-3) + x = torch.randn(1, small_config.num_channels, + small_config.image_size, small_config.image_size) + t = torch.zeros(1, dtype=torch.int64) + z = torch.randn(1, small_config.nz) + target = torch.full_like(x, 0.5) + + losses = [] + for _ in range(15): + opt.zero_grad() + loss = torch.nn.functional.mse_loss(net(x, t, z), target) + loss.backward() + opt.step() + losses.append(loss.item()) + + assert losses[-1] < losses[0], "loss did not move: {} -> {}".format(losses[0], losses[-1]) + + +@pytest.mark.parametrize("batch", [1, 2, 3, 4, 5, 6, 7, 8]) +def test_discriminator_large_accepts_any_batch_size(batch): + """The minibatch-stddev grouping must divide the batch, whatever its size.""" + net = Discriminator_large(nc=2, ngf=8, t_emb_dim=16, act=nn.LeakyReLU(0.2)) + out = net(torch.randn(batch, 1, 128, 128), + torch.randint(0, 4, (batch,)), + torch.randn(batch, 1, 128, 128)) + assert out.shape == (batch, 1) + + +@pytest.mark.parametrize("batch", [1, 3, 5, 6]) +def test_discriminator_small_accepts_any_batch_size(batch): + net = Discriminator_small(nc=2, ngf=8, t_emb_dim=16, act=nn.LeakyReLU(0.2)) + out = net(torch.randn(batch, 1, 32, 32), + torch.randint(0, 4, (batch,)), + torch.randn(batch, 1, 32, 32)) + assert out.shape == (batch, 1) + + +def test_discriminator_large_backward(): + net = Discriminator_large(nc=2, ngf=8, t_emb_dim=16, act=nn.LeakyReLU(0.2)) + out = net(torch.randn(2, 1, 128, 128), torch.randint(0, 4, (2,)), + torch.randn(2, 1, 128, 128)) + out.sum().backward() + assert any(p.grad is not None and torch.isfinite(p.grad).all() + for p in net.parameters()) + + +def test_define_g_is_not_wrapped_in_a_parallel_module(): + """init_net must leave parallelism to DDP; a nested wrapper renames keys.""" + net = generator_resnet.define_G(netG='resnet_6blocks', gpu_ids=[]) + assert not isinstance(net, nn.DataParallel) + assert all(not k.startswith('module.') for k in net.state_dict()) + + +def test_translation_networks_preserve_spatial_shape(): + gen = generator_resnet.define_G(netG='resnet_6blocks', gpu_ids=[]) + x = torch.randn(2, 1, 64, 64) + out = gen(x) + assert out.shape == x.shape + + disc = generator_resnet.define_D(gpu_ids=[]) + # PatchGAN returns a map of patch scores rather than one scalar + assert disc(out).ndim == 4 diff --git a/tests/test_ops.py b/tests/test_ops.py new file mode 100644 index 00000000..79aa4725 --- /dev/null +++ b/tests/test_ops.py @@ -0,0 +1,72 @@ +"""The CUDA kernels and their pure-PyTorch fallbacks must agree. + +utils/op builds its extensions lazily and falls back when the toolchain is +missing, so these tests only run where the kernels are actually available. +""" +import sys + +import pytest +import torch + +import utils.op # noqa: F401 (populates sys.modules with the submodules) + +upfirdn2d_module = sys.modules['utils.op.upfirdn2d'] +fused_act_module = sys.modules['utils.op.fused_act'] + +needs_kernels = pytest.mark.skipif( + not torch.cuda.is_available() + or upfirdn2d_module.upfirdn2d_op is None + or fused_act_module.fused is None, + reason="CUDA extensions are not built on this machine", +) + + +def fir_kernel(device): + k = torch.tensor([1., 3., 3., 1.]) + k = torch.outer(k, k) + return (k / k.sum()).to(device) + + +def test_fallback_is_selected_when_the_extension_is_missing(monkeypatch): + """A CPU tensor must never reach the CUDA path.""" + monkeypatch.setattr(upfirdn2d_module, 'UpFirDn2d', None) + out = upfirdn2d_module.upfirdn2d(torch.randn(1, 1, 8, 8), + fir_kernel('cpu'), up=1, down=1, pad=(1, 1)) + assert out.shape == (1, 1, 7, 7) + + +@needs_kernels +@pytest.mark.parametrize("up,down,pad", [(2, 1, (2, 1)), (1, 2, (1, 1)), (1, 1, (1, 1))]) +def test_upfirdn2d_matches_the_native_implementation(up, down, pad): + torch.manual_seed(0) + x = torch.randn(2, 3, 16, 16, device='cuda') + k = fir_kernel('cuda') + cuda_out = upfirdn2d_module.UpFirDn2d.apply( + x, k, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])) + native_out = upfirdn2d_module.upfirdn2d_native( + x, k, up, up, down, down, pad[0], pad[1], pad[0], pad[1]) + assert cuda_out.shape == native_out.shape + assert torch.allclose(cuda_out, native_out, atol=1e-6) + + +@needs_kernels +def test_upfirdn2d_gradients_match_the_native_implementation(): + torch.manual_seed(0) + k = fir_kernel('cuda') + a = torch.randn(2, 3, 16, 16, device='cuda', requires_grad=True) + b = a.detach().clone().requires_grad_(True) + upfirdn2d_module.UpFirDn2d.apply(a, k, (2, 2), (1, 1), (2, 1, 2, 1)).square().sum().backward() + upfirdn2d_module.upfirdn2d_native(b, k, 2, 2, 1, 1, 2, 1, 2, 1).square().sum().backward() + assert torch.allclose(a.grad, b.grad, atol=1e-5) + + +@needs_kernels +def test_fused_leaky_relu_matches_the_native_implementation(): + torch.manual_seed(0) + x = torch.randn(4, 8, 16, 16, device='cuda') + bias = torch.randn(8, device='cuda') + scale = 2 ** 0.5 + cuda_out = fused_act_module.FusedLeakyReLUFunction.apply(x, bias, 0.2, scale) + native_out = torch.nn.functional.leaky_relu( + x + bias.view(1, -1, 1, 1), negative_slope=0.2) * scale + assert torch.allclose(cuda_out, native_out, atol=1e-6) From e1dd4a7039696ee68dd18b55a1432861604ad8bd Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:10:03 +0300 Subject: [PATCH 16/21] chore: add requirements.txt, ignore .DS_Store and local envs The repository listed its dependencies only as prose in the README and shipped no requirements file, so there was nothing to install from. Record the actual runtime set, and note why torch>=1.13 is the floor: train.py now passes weights_only to torch.load, which that release introduced. ninja is listed separately because it is only needed to build the fused CUDA kernels; without it SynDiff falls back to the native implementations. Also drop the committed .DS_Store and extend .gitignore to cover it along with local virtualenvs and pytest caches. --- .DS_Store | Bin 6148 -> 0 bytes .gitignore | 10 +++++++++- requirements.txt | 18 ++++++++++++++++++ 3 files changed, 27 insertions(+), 1 deletion(-) delete mode 100644 .DS_Store create mode 100644 requirements.txt diff --git a/.DS_Store b/.DS_Store deleted file mode 100644 index f51ca32b3e7b4d153ef5c37bfdd695269bcb02da..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6148 zcmeHKO>fgM7=GQA%~3U}2av|i8rMoRL%H>>)7dUu*=EMqfpN&?NW+y*6<2XT8&bTeujFPEUMCujAuF_OUIzRd6gBZ zTnyyANGqKinhV_HUevw2Tpm1l*c1NY!Kx>gzTbb`6OW!8u2vqubMOAsxt3GFDezAzz}mxbIKY(b z-nub4wbu&xJGe5%YYcu*K}B!H;Hs^753UVyp(?<@Vq*{ynEMgXGPuGi@JAK+21DC+ A`2YX_ diff --git a/.gitignore b/.gitignore index d36252ad..7cdf2146 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,12 @@ *.pyc *pycache* saved_info/* -saved_info \ No newline at end of file +saved_info + +# local environments and scratch output +.venv/ +venv/ +.pytest_cache/ + +# macOS +.DS_Store diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 00000000..79ec4ac3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +# Runtime dependencies. +# +# torch>=1.13 is required because train.py passes weights_only to torch.load; +# the argument was introduced in that release. PyTorch 2.6 flipped its default +# to True, which is why passing it explicitly matters. +torch>=1.13 +torchvision>=0.14 +numpy +h5py +scikit-image + +# Optional: needed only to JIT-build the fused CUDA kernels in utils/op. +# Without it (or without a matching CUDA toolchain) SynDiff warns once and +# falls back to the equivalent pure-PyTorch implementations. +ninja + +# Tests +pytest From 0009353330d20b343f3d90b4c98922f5a0deaace Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:10:42 +0300 Subject: [PATCH 17/21] docs: correct the dependency, dataset and inference notes Several statements in the README did not match the code: - The dependency list omitted numpy, h5py and scikit-image, all of which are imported at module scope, and pinned torch>=1.7.1, which is below the floor the resume path now needs. ninja and the CUDA toolchain moved to an optional section, since the fused kernels fall back to native PyTorch. - The dataset section did not mention that each .mat file must hold a variable called data_fs, nor that volumes are padded to 256x256 and rescaled to [-1, 1] on load, so neither dimension may exceed 256. - The sample data was described as ready to use. The folder actually holds two raw 25-slice volumes named T1.mat/T2.mat, which do not match the data__.mat layout the loader expects; say so. - Nothing described what --num_process_per_node changes, where test.py writes its output, or that the crop applied before saving is configurable. Also document how to run the test suite. --- README.md | 65 ++++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 60 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index f271c74e..02160b32 100644 --- a/README.md +++ b/README.md @@ -14,18 +14,36 @@ Muzaffer Özbey*, Onat Dalmaz*, Salman UH Dar, Hasan A Bedel, Şaban Özturk, Al ``` python>=3.6.9 -torch>=1.7.1 -torchvision>=0.8.2 -cuda=>11.2 +torch>=1.13 +torchvision>=0.14 +numpy +h5py +scikit-image +``` + +`torch>=1.13` is required because the resume path passes `weights_only` to +`torch.load`; PyTorch 2.6 later flipped that argument's default to `True`, +which is why it is now passed explicitly. + +### Optional: fused CUDA kernels +`utils/op` ships hand-written CUDA kernels that are JIT-compiled on first +import. Building them needs a CUDA toolchain and: + +``` +cuda>=11.2 ninja python3.x-dev (apt install, x should match your python3 version, ex: 3.8) ``` +If any of these is missing, SynDiff warns once and falls back to equivalent +pure-PyTorch implementations, so the code also runs on a CPU-only install. + ## Installation - Clone this repo: ```bash git clone https://github.com/icon-lab/SynDiff cd SynDiff +pip install -r requirements.txt ``` ## Dataset @@ -43,9 +61,19 @@ input_path/ ├── data_test_contrast2.mat ``` -where .mat files has shape of (#images, width, height) and image values are between 0 and 1.0. +where the `contrast1`/`contrast2` parts of the file names are the values passed +to `--contrast1` and `--contrast2`. + +Each `.mat` file is an HDF5 file holding a single variable named `data_fs` of +shape `(#images, width, height)`, with image values in roughly `[0, 1]`. +Volumes are zero-padded out to 256x256 on load and rescaled to `[-1, 1]`, so +neither dimension may exceed 256. + ### Sample Data -Sample toy data can also found under 'SynDiff_sample_data' folder of the repository. +Sample toy data can be found under the `SynDiff_sample_data` folder. Note that +those two files are raw volumes (`T1.mat`, `T2.mat`, 25 slices each) rather +than a ready-made split -- to run the commands below, split them into train / +val / test parts and name the parts as shown above. @@ -57,6 +85,11 @@ Sample toy data can also found under 'SynDiff_sample_data' folder of the reposit python3 train.py --image_size 256 --exp exp_syndiff --num_channels 2 --num_channels_dae 64 --ch_mult 1 1 2 2 4 4 --num_timesteps 4 --num_res_blocks 2 --batch_size 1 --contrast1 T1 --contrast2 T2 --num_epoch 500 --ngf 64 --embedding_type positional --use_ema --ema_decay 0.999 --r1_gamma 1. --z_emb_dim 256 --lr_d 1e-4 --lr_g 1.6e-4 --lazy_reg 10 --num_process_per_node 1 --save_content --local_rank 0 --input_path /input/path/for/data --output_path /output/for/results ``` +`--num_process_per_node` controls the number of processes. With more than one +a NCCL process group is set up and the networks are wrapped in +`DistributedDataParallel`; with a single process neither is used, and the run +falls back to CPU when no GPU is visible. +
## Pretrained Models @@ -70,6 +103,28 @@ We have released pretrained diffusive generators for [T1->PD and PD->T1](https:/ python test.py --image_size 256 --exp exp_syndiff --num_channels 2 --num_channels_dae 64 --ch_mult 1 1 2 2 4 4 --num_timesteps 4 --num_res_blocks 2 --batch_size 1 --embedding_type positional --z_emb_dim 256 --contrast1 T1 --contrast2 T2 --which_epoch 50 --gpu_chose 0 --input_path /input/path/for/data --output_path /output/for/results ``` +Synthesised images are written to +`output_path/exp/generated_samples/epoch_/`, both as JPEGs and +collected into `im_syn.mat`. Before saving, each image is cropped back from +the padded 256x256 grid; `--crop_h` and `--crop_w` set that size and default +to `256 152`, the slice geometry used in the paper. Set them to your own +slice size for other datasets. + +
+ +## Tests + +A CPU test suite covers the diffusion coefficients, network shapes and +gradients, dataset loading and checkpoint handling: + +``` +pip install -r requirements.txt +python -m pytest tests/ +``` + +Tests that compare the fused CUDA kernels against their pure-PyTorch +fallbacks are skipped automatically when the extensions cannot be built. +

From 39252bf3a6fed95767c6c032adf176b448c587bc Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:10:57 +0300 Subject: [PATCH 18/21] chore: remove two unreferenced modules backbones/im2im.py holds a second copy of define_G/init_net that nothing imports; the live copy is in backbones/generator_resnet.py. Keeping both means fixes have to be applied twice, and the one just made to init_net (dropping the DataParallel wrapper) would silently not apply here. utils/utils.py holds restore_checkpoint/save_checkpoint, which nothing calls, and imports tensorflow at module scope for a single tf.io.gfile existence check. That makes TensorFlow look like a dependency of a PyTorch project. Neither file is referenced from train.py, test.py, dataset.py or backbones. --- backbones/im2im.py | 182 --------------------------------------------- utils/utils.py | 30 -------- 2 files changed, 212 deletions(-) delete mode 100644 backbones/im2im.py delete mode 100644 utils/utils.py diff --git a/backbones/im2im.py b/backbones/im2im.py deleted file mode 100644 index a44fc82b..00000000 --- a/backbones/im2im.py +++ /dev/null @@ -1,182 +0,0 @@ -# --------------------------------------------------------------- -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. -# --------------------------------------------------------------- - -# coding=utf-8 -# Copyright 2020 The Google Research Authors. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# pylint: skip-file -''' Codes adapted from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ncsnpp.py -''' - -from . import utils, layers, layerspp, dense_layer -import torch.nn as nn -import functools -import torch -import numpy as np - -def define_G(input_nc=1, output_nc=1, ngf=64, netG='resnet_9blocks', norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]): - """Create a generator - Parameters: - input_nc (int) -- the number of channels in input images - output_nc (int) -- the number of channels in output images - ngf (int) -- the number of filters in the last conv layer - netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128 - norm (str) -- the name of normalization layers used in the network: batch | instance | none - use_dropout (bool) -- if use dropout layers. - init_type (str) -- the name of our initialization method. - init_gain (float) -- scaling factor for normal, xavier and orthogonal. - gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 - Returns a generator - Our current implementation provides two types of generators: - U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images) - The original U-Net paper: https://arxiv.org/abs/1505.04597 - Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks) - Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations. - We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style). - The generator has been initialized by . It uses RELU for non-linearity. - """ - net = None - norm_layer = get_norm_layer(norm_type=norm) - - if netG == 'resnet_9blocks': - net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9) - elif netG == 'resnet_6blocks': - net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6) - elif netG == 'unet_128': - net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout) - elif netG == 'unet_256': - net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout) - else: - raise NotImplementedError('Generator model name [%s] is not recognized' % netG) - return init_net(net, init_type, init_gain, gpu_ids) - - - -class ResnetGenerator(nn.Module): - """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations. - We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style) - """ - - def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=9, padding_type='reflect'): - """Construct a Resnet-based generator - Parameters: - input_nc (int) -- the number of channels in input images - output_nc (int) -- the number of channels in output images - ngf (int) -- the number of filters in the last conv layer - norm_layer -- normalization layer - use_dropout (bool) -- if use dropout layers - n_blocks (int) -- the number of ResNet blocks - padding_type (str) -- the name of padding layer in conv layers: reflect | replicate | zero - """ - assert(n_blocks >= 0) - super(ResnetGenerator, self).__init__() - if type(norm_layer) == functools.partial: - use_bias = norm_layer.func == nn.InstanceNorm2d - else: - use_bias = norm_layer == nn.InstanceNorm2d - - model = [nn.ReflectionPad2d(3), - nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), - norm_layer(ngf), - nn.ReLU(True)] - - n_downsampling = 2 - for i in range(n_downsampling): # add downsampling layers - mult = 2 ** i - model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias), - norm_layer(ngf * mult * 2), - nn.ReLU(True)] - - mult = 2 ** n_downsampling - for i in range(n_blocks): # add ResNet blocks - - model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)] - - for i in range(n_downsampling): # add upsampling layers - mult = 2 ** (n_downsampling - i) - model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), - kernel_size=3, stride=2, - padding=1, output_padding=1, - bias=use_bias), - norm_layer(int(ngf * mult / 2)), - nn.ReLU(True)] - model += [nn.ReflectionPad2d(3)] - model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)] - model += [nn.Tanh()] - - self.model = nn.Sequential(*model) - - def forward(self, input): - """Standard forward""" - return self.model(input) - - -class ResnetBlock(nn.Module): - """Define a Resnet block""" - - def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias): - """Initialize the Resnet block - A resnet block is a conv block with skip connections - We construct a conv block with build_conv_block function, - and implement skip connections in function. - Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf - """ - super(ResnetBlock, self).__init__() - self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias) - - def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias): - """Construct a convolutional block. - Parameters: - dim (int) -- the number of channels in the conv layer. - padding_type (str) -- the name of padding layer: reflect | replicate | zero - norm_layer -- normalization layer - use_dropout (bool) -- if use dropout layers. - use_bias (bool) -- if the conv layer uses bias or not - Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU)) - """ - conv_block = [] - p = 0 - if padding_type == 'reflect': - conv_block += [nn.ReflectionPad2d(1)] - elif padding_type == 'replicate': - conv_block += [nn.ReplicationPad2d(1)] - elif padding_type == 'zero': - p = 1 - else: - raise NotImplementedError('padding [%s] is not implemented' % padding_type) - - conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)] - if use_dropout: - conv_block += [nn.Dropout(0.5)] - - p = 0 - if padding_type == 'reflect': - conv_block += [nn.ReflectionPad2d(1)] - elif padding_type == 'replicate': - conv_block += [nn.ReplicationPad2d(1)] - elif padding_type == 'zero': - p = 1 - else: - raise NotImplementedError('padding [%s] is not implemented' % padding_type) - conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)] - - return nn.Sequential(*conv_block) - - def forward(self, x): - """Forward function (with skip connections)""" - out = x + self.conv_block(x) # add skip connections - return out - diff --git a/utils/utils.py b/utils/utils.py deleted file mode 100644 index 19236035..00000000 --- a/utils/utils.py +++ /dev/null @@ -1,30 +0,0 @@ - -import torch -import tensorflow as tf -import os -import logging - - -def restore_checkpoint(ckpt_dir, state, device): - if not tf.io.gfile.exists(ckpt_dir): - tf.io.gfile.makedirs(os.path.dirname(ckpt_dir)) - logging.warning(f"No checkpoint found at {ckpt_dir}. " - f"Returned the same state as input") - return state - else: - loaded_state = torch.load(ckpt_dir, map_location=device) - state['optimizer'].load_state_dict(loaded_state['optimizer']) - state['model'].load_state_dict(loaded_state['model'], strict=False) - state['ema'].load_state_dict(loaded_state['ema']) - state['step'] = loaded_state['step'] - return state - - -def save_checkpoint(ckpt_dir, state): - saved_state = { - 'optimizer': state['optimizer'].state_dict(), - 'model': state['model'].state_dict(), - 'ema': state['ema'].state_dict(), - 'step': state['step'] - } - torch.save(saved_state, ckpt_dir) \ No newline at end of file From 6416389311b4d09944ba8397236a960b9e50a25a Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:19:21 +0300 Subject: [PATCH 19/21] refactor: move the shared diffusion code into diffusion.py train.py and test.py each carried their own copy of var_func_vp, var_func_geometric, extract, get_time_schedule, get_sigma_schedule, Posterior_Coefficients, sample_posterior and sample_from_model -- about 120 duplicated lines defining the sampler that both the training loop and inference depend on. The two copies are currently byte-identical, so any change to the diffusion process has to be made twice and silently produces a train/test mismatch if it is not. Move them into diffusion.py, together with Diffusion_Coefficients, q_sample and q_sample_pairs, which only train.py had. The definitions are moved verbatim; both entry points now import from the single copy. Behaviour is unchanged: the extracted sources were checked to be identical between the two files before merging. --- diffusion.py | 169 ++++++++++++++++++++++++++++++++++++++++ test.py | 116 +-------------------------- tests/test_diffusion.py | 6 +- train.py | 153 +----------------------------------- 4 files changed, 177 insertions(+), 267 deletions(-) create mode 100644 diffusion.py diff --git a/diffusion.py b/diffusion.py new file mode 100644 index 00000000..5bdba795 --- /dev/null +++ b/diffusion.py @@ -0,0 +1,169 @@ +"""Diffusion coefficients and sampling shared by training and inference. + +Extracted verbatim from train.py/test.py, which each carried their own copy of +these definitions. +""" + +import numpy as np +import torch + + +def var_func_vp(t, beta_min, beta_max): + log_mean_coeff = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min + var = 1. - torch.exp(2. * log_mean_coeff) + return var + + +def var_func_geometric(t, beta_min, beta_max): + return beta_min * ((beta_max / beta_min) ** t) + + +def extract(input, t, shape): + out = torch.gather(input, 0, t) + reshape = [shape[0]] + [1] * (len(shape) - 1) + out = out.reshape(*reshape) + + return out + + +def get_time_schedule(args, device): + n_timestep = args.num_timesteps + eps_small = 1e-3 + t = np.arange(0, n_timestep + 1, dtype=np.float64) + t = t / n_timestep + t = torch.from_numpy(t) * (1. - eps_small) + eps_small + return t.to(device) + + +def get_sigma_schedule(args, device): + n_timestep = args.num_timesteps + beta_min = args.beta_min + beta_max = args.beta_max + eps_small = 1e-3 + + t = np.arange(0, n_timestep + 1, dtype=np.float64) + t = t / n_timestep + t = torch.from_numpy(t) * (1. - eps_small) + eps_small + + if args.use_geometric: + var = var_func_geometric(t, beta_min, beta_max) + else: + var = var_func_vp(t, beta_min, beta_max) + alpha_bars = 1.0 - var + betas = 1 - alpha_bars[1:] / alpha_bars[:-1] + + first = torch.tensor(1e-8) + betas = torch.cat((first[None], betas)).to(device) + betas = betas.type(torch.float32) + sigmas = betas**0.5 + a_s = torch.sqrt(1-betas) + return sigmas, a_s, betas + + +class Diffusion_Coefficients(): + def __init__(self, args, device): + + self.sigmas, self.a_s, _ = get_sigma_schedule(args, device=device) + self.a_s_cum = np.cumprod(self.a_s.cpu()) + self.sigmas_cum = np.sqrt(1 - self.a_s_cum ** 2) + self.a_s_prev = self.a_s.clone() + self.a_s_prev[-1] = 1 + + self.a_s_cum = self.a_s_cum.to(device) + self.sigmas_cum = self.sigmas_cum.to(device) + self.a_s_prev = self.a_s_prev.to(device) + + +def q_sample(coeff, x_start, t, *, noise=None): + """ + Diffuse the data (t == 0 means diffused for t step) + """ + if noise is None: + noise = torch.randn_like(x_start) + + x_t = extract(coeff.a_s_cum, t, x_start.shape) * x_start + \ + extract(coeff.sigmas_cum, t, x_start.shape) * noise + + return x_t + + +def q_sample_pairs(coeff, x_start, t): + """ + Generate a pair of disturbed images for training + :param x_start: x_0 + :param t: time step t + :return: x_t, x_{t+1} + """ + noise = torch.randn_like(x_start) + x_t = q_sample(coeff, x_start, t) + x_t_plus_one = extract(coeff.a_s, t+1, x_start.shape) * x_t + \ + extract(coeff.sigmas, t+1, x_start.shape) * noise + + return x_t, x_t_plus_one + + +class Posterior_Coefficients(): + def __init__(self, args, device): + + _, _, self.betas = get_sigma_schedule(args, device=device) + + #we don't need the zeros + self.betas = self.betas.type(torch.float32)[1:] + + self.alphas = 1 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, 0) + self.alphas_cumprod_prev = torch.cat( + (torch.tensor([1.], dtype=torch.float32,device=device), self.alphas_cumprod[:-1]), 0 + ) + self.posterior_variance = self.betas * (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) + + self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod) + self.sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod) + self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1 / self.alphas_cumprod - 1) + + self.posterior_mean_coef1 = (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1 - self.alphas_cumprod)) + self.posterior_mean_coef2 = ((1 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1 - self.alphas_cumprod)) + + self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min=1e-20)) + + +def sample_posterior(coefficients, x_0,x_t, t): + + def q_posterior(x_0, x_t, t): + mean = ( + extract(coefficients.posterior_mean_coef1, t, x_t.shape) * x_0 + + extract(coefficients.posterior_mean_coef2, t, x_t.shape) * x_t + ) + var = extract(coefficients.posterior_variance, t, x_t.shape) + log_var_clipped = extract(coefficients.posterior_log_variance_clipped, t, x_t.shape) + return mean, var, log_var_clipped + + + def p_sample(x_0, x_t, t): + mean, _, log_var = q_posterior(x_0, x_t, t) + + noise = torch.randn_like(x_t) + + nonzero_mask = (1 - (t == 0).type(torch.float32)) + + return mean + nonzero_mask[:,None,None,None] * torch.exp(0.5 * log_var) * noise + + sample_x_pos = p_sample(x_0, x_t, t) + + return sample_x_pos + + +def sample_from_model(coefficients, generator, n_time, x_init, T, opt): + x = x_init[:,[0],:] + source = x_init[:,[1],:] + with torch.no_grad(): + for i in reversed(range(n_time)): + t = torch.full((x.size(0),), i, dtype=torch.int64).to(x.device) + + t_time = t + latent_z = torch.randn(x.size(0), opt.nz, device=x.device)#.to(x.device) + x_0 = generator(torch.cat((x,source),axis=1), t_time, latent_z) + x_new = sample_posterior(coefficients, x_0[:,[0],:], x, t) + x = x_new.detach() + + return x diff --git a/test.py b/test.py index 593ed422..ee58641c 100644 --- a/test.py +++ b/test.py @@ -7,6 +7,8 @@ import torch.optim as optim import torchvision from backbones.ncsnpp_generator_adagn import NCSNpp +from diffusion import (Posterior_Coefficients, get_time_schedule, + sample_from_model) from dataset import CreateDatasetSynthesis import torch.nn.functional as F @@ -19,120 +21,6 @@ def psnr(img1, img2): mse = torch.mean((img1 - img2) ** 2) return 20 * torch.log10(img1.max() / torch.sqrt(mse)) - -#%% Diffusion coefficients -def var_func_vp(t, beta_min, beta_max): - log_mean_coeff = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min - var = 1. - torch.exp(2. * log_mean_coeff) - return var - -def var_func_geometric(t, beta_min, beta_max): - return beta_min * ((beta_max / beta_min) ** t) - -def extract(input, t, shape): - out = torch.gather(input, 0, t) - reshape = [shape[0]] + [1] * (len(shape) - 1) - out = out.reshape(*reshape) - - return out - -def get_time_schedule(args, device): - n_timestep = args.num_timesteps - eps_small = 1e-3 - t = np.arange(0, n_timestep + 1, dtype=np.float64) - t = t / n_timestep - t = torch.from_numpy(t) * (1. - eps_small) + eps_small - return t.to(device) - -def get_sigma_schedule(args, device): - n_timestep = args.num_timesteps - beta_min = args.beta_min - beta_max = args.beta_max - eps_small = 1e-3 - - t = np.arange(0, n_timestep + 1, dtype=np.float64) - t = t / n_timestep - t = torch.from_numpy(t) * (1. - eps_small) + eps_small - - if args.use_geometric: - var = var_func_geometric(t, beta_min, beta_max) - else: - var = var_func_vp(t, beta_min, beta_max) - alpha_bars = 1.0 - var - betas = 1 - alpha_bars[1:] / alpha_bars[:-1] - - first = torch.tensor(1e-8) - betas = torch.cat((first[None], betas)).to(device) - betas = betas.type(torch.float32) - sigmas = betas**0.5 - a_s = torch.sqrt(1-betas) - return sigmas, a_s, betas - -#%% posterior sampling -class Posterior_Coefficients(): - def __init__(self, args, device): - - _, _, self.betas = get_sigma_schedule(args, device=device) - - #we don't need the zeros - self.betas = self.betas.type(torch.float32)[1:] - - self.alphas = 1 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, 0) - self.alphas_cumprod_prev = torch.cat( - (torch.tensor([1.], dtype=torch.float32,device=device), self.alphas_cumprod[:-1]), 0 - ) - self.posterior_variance = self.betas * (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) - - self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod) - self.sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod) - self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1 / self.alphas_cumprod - 1) - - self.posterior_mean_coef1 = (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1 - self.alphas_cumprod)) - self.posterior_mean_coef2 = ((1 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1 - self.alphas_cumprod)) - - self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min=1e-20)) - -def sample_posterior(coefficients, x_0,x_t, t): - - def q_posterior(x_0, x_t, t): - mean = ( - extract(coefficients.posterior_mean_coef1, t, x_t.shape) * x_0 - + extract(coefficients.posterior_mean_coef2, t, x_t.shape) * x_t - ) - var = extract(coefficients.posterior_variance, t, x_t.shape) - log_var_clipped = extract(coefficients.posterior_log_variance_clipped, t, x_t.shape) - return mean, var, log_var_clipped - - - def p_sample(x_0, x_t, t): - mean, _, log_var = q_posterior(x_0, x_t, t) - - noise = torch.randn_like(x_t) - - nonzero_mask = (1 - (t == 0).type(torch.float32)) - - return mean + nonzero_mask[:,None,None,None] * torch.exp(0.5 * log_var) * noise - - sample_x_pos = p_sample(x_0, x_t, t) - - return sample_x_pos - -def sample_from_model(coefficients, generator, n_time, x_init, T, opt): - x = x_init[:,[0],:] - source = x_init[:,[1],:] - with torch.no_grad(): - for i in reversed(range(n_time)): - t = torch.full((x.size(0),), i, dtype=torch.int64).to(x.device) - - t_time = t - latent_z = torch.randn(x.size(0), opt.nz, device=x.device)#.to(x.device) - x_0 = generator(torch.cat((x,source),axis=1), t_time, latent_z) - x_new = sample_posterior(coefficients, x_0[:,[0],:], x, t) - x = x_new.detach() - - return x - def load_checkpoint(checkpoint_dir, netG, name_of_network, epoch,device = 'cuda:0'): checkpoint_file = checkpoint_dir.format(name_of_network, epoch) diff --git a/tests/test_diffusion.py b/tests/test_diffusion.py index f5e6b638..1fcb3fba 100644 --- a/tests/test_diffusion.py +++ b/tests/test_diffusion.py @@ -3,9 +3,9 @@ import pytest import torch -from train import (Diffusion_Coefficients, Posterior_Coefficients, extract, - get_sigma_schedule, get_time_schedule, q_sample_pairs, - sample_posterior) +from diffusion import (Diffusion_Coefficients, Posterior_Coefficients, extract, + get_sigma_schedule, get_time_schedule, q_sample_pairs, + sample_posterior) def test_sigma_schedule_shapes_and_bounds(small_config): diff --git a/train.py b/train.py index 39f7be32..0e53ce77 100644 --- a/train.py +++ b/train.py @@ -13,6 +13,9 @@ import torchvision import torchvision.transforms as transforms +from diffusion import (Diffusion_Coefficients, Posterior_Coefficients, + get_time_schedule, q_sample_pairs, sample_from_model, + sample_posterior) from dataset import CreateDatasetSynthesis from torch.multiprocessing import Process @@ -72,156 +75,6 @@ def load_model_state(model, state_dict): target.load_state_dict(strip_module_prefix(state_dict)) -#%% Diffusion coefficients -def var_func_vp(t, beta_min, beta_max): - log_mean_coeff = -0.25 * t ** 2 * (beta_max - beta_min) - 0.5 * t * beta_min - var = 1. - torch.exp(2. * log_mean_coeff) - return var - -def var_func_geometric(t, beta_min, beta_max): - return beta_min * ((beta_max / beta_min) ** t) - -def extract(input, t, shape): - out = torch.gather(input, 0, t) - reshape = [shape[0]] + [1] * (len(shape) - 1) - out = out.reshape(*reshape) - - return out - -def get_time_schedule(args, device): - n_timestep = args.num_timesteps - eps_small = 1e-3 - t = np.arange(0, n_timestep + 1, dtype=np.float64) - t = t / n_timestep - t = torch.from_numpy(t) * (1. - eps_small) + eps_small - return t.to(device) - -def get_sigma_schedule(args, device): - n_timestep = args.num_timesteps - beta_min = args.beta_min - beta_max = args.beta_max - eps_small = 1e-3 - - t = np.arange(0, n_timestep + 1, dtype=np.float64) - t = t / n_timestep - t = torch.from_numpy(t) * (1. - eps_small) + eps_small - - if args.use_geometric: - var = var_func_geometric(t, beta_min, beta_max) - else: - var = var_func_vp(t, beta_min, beta_max) - alpha_bars = 1.0 - var - betas = 1 - alpha_bars[1:] / alpha_bars[:-1] - - first = torch.tensor(1e-8) - betas = torch.cat((first[None], betas)).to(device) - betas = betas.type(torch.float32) - sigmas = betas**0.5 - a_s = torch.sqrt(1-betas) - return sigmas, a_s, betas - -class Diffusion_Coefficients(): - def __init__(self, args, device): - - self.sigmas, self.a_s, _ = get_sigma_schedule(args, device=device) - self.a_s_cum = np.cumprod(self.a_s.cpu()) - self.sigmas_cum = np.sqrt(1 - self.a_s_cum ** 2) - self.a_s_prev = self.a_s.clone() - self.a_s_prev[-1] = 1 - - self.a_s_cum = self.a_s_cum.to(device) - self.sigmas_cum = self.sigmas_cum.to(device) - self.a_s_prev = self.a_s_prev.to(device) - -def q_sample(coeff, x_start, t, *, noise=None): - """ - Diffuse the data (t == 0 means diffused for t step) - """ - if noise is None: - noise = torch.randn_like(x_start) - - x_t = extract(coeff.a_s_cum, t, x_start.shape) * x_start + \ - extract(coeff.sigmas_cum, t, x_start.shape) * noise - - return x_t - -def q_sample_pairs(coeff, x_start, t): - """ - Generate a pair of disturbed images for training - :param x_start: x_0 - :param t: time step t - :return: x_t, x_{t+1} - """ - noise = torch.randn_like(x_start) - x_t = q_sample(coeff, x_start, t) - x_t_plus_one = extract(coeff.a_s, t+1, x_start.shape) * x_t + \ - extract(coeff.sigmas, t+1, x_start.shape) * noise - - return x_t, x_t_plus_one -#%% posterior sampling -class Posterior_Coefficients(): - def __init__(self, args, device): - - _, _, self.betas = get_sigma_schedule(args, device=device) - - #we don't need the zeros - self.betas = self.betas.type(torch.float32)[1:] - - self.alphas = 1 - self.betas - self.alphas_cumprod = torch.cumprod(self.alphas, 0) - self.alphas_cumprod_prev = torch.cat( - (torch.tensor([1.], dtype=torch.float32,device=device), self.alphas_cumprod[:-1]), 0 - ) - self.posterior_variance = self.betas * (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) - - self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod) - self.sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod) - self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1 / self.alphas_cumprod - 1) - - self.posterior_mean_coef1 = (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1 - self.alphas_cumprod)) - self.posterior_mean_coef2 = ((1 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1 - self.alphas_cumprod)) - - self.posterior_log_variance_clipped = torch.log(self.posterior_variance.clamp(min=1e-20)) - -def sample_posterior(coefficients, x_0,x_t, t): - - def q_posterior(x_0, x_t, t): - mean = ( - extract(coefficients.posterior_mean_coef1, t, x_t.shape) * x_0 - + extract(coefficients.posterior_mean_coef2, t, x_t.shape) * x_t - ) - var = extract(coefficients.posterior_variance, t, x_t.shape) - log_var_clipped = extract(coefficients.posterior_log_variance_clipped, t, x_t.shape) - return mean, var, log_var_clipped - - - def p_sample(x_0, x_t, t): - mean, _, log_var = q_posterior(x_0, x_t, t) - - noise = torch.randn_like(x_t) - - nonzero_mask = (1 - (t == 0).type(torch.float32)) - - return mean + nonzero_mask[:,None,None,None] * torch.exp(0.5 * log_var) * noise - - sample_x_pos = p_sample(x_0, x_t, t) - - return sample_x_pos - -def sample_from_model(coefficients, generator, n_time, x_init, T, opt): - x = x_init[:,[0],:] - source = x_init[:,[1],:] - with torch.no_grad(): - for i in reversed(range(n_time)): - t = torch.full((x.size(0),), i, dtype=torch.int64).to(x.device) - - t_time = t - latent_z = torch.randn(x.size(0), opt.nz, device=x.device)#.to(x.device) - x_0 = generator(torch.cat((x,source),axis=1), t_time, latent_z) - x_new = sample_posterior(coefficients, x_0[:,[0],:], x, t) - x = x_new.detach() - - return x #%% def train_syndiff(rank, gpu, args): From f26c20797b81acfd8caf3a8910c9f0c0413e9615 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:19:27 +0300 Subject: [PATCH 20/21] fix: honour --ngf and --image_size outside the diffusive networks Two command-line arguments only reached part of the code they name. --ngf was passed to the two diffusive discriminators, but the translation networks and the cycle discriminators were built through define_G/define_D without it, so they always used the functions' own default of 64. Anyone training with a different --ngf silently got a model whose halves disagreed, with no way to scale the translation path at all. --image_size was passed to NCSNpp, but CreateDatasetSynthesis padded every volume to a hard-coded 256x256. Requesting any other size gave a network configured for one resolution and data at another, and SynDiff could not be used on a 256-grid-incompatible dataset without editing dataset.py. Thread both through. The published commands pass --ngf 64 and --image_size 256, which are exactly the previous hard-coded values, so they produce the same networks and the same data as before. --- dataset.py | 9 +++++---- test.py | 2 +- tests/test_dataset.py | 23 +++++++++++++++++++++++ tests/test_models.py | 15 +++++++++++++++ train.py | 12 ++++++------ 5 files changed, 50 insertions(+), 11 deletions(-) diff --git a/dataset.py b/dataset.py index b1762029..ce8f63b7 100644 --- a/dataset.py +++ b/dataset.py @@ -5,13 +5,14 @@ import random -def CreateDatasetSynthesis(phase, input_path, contrast1 = 'T1', contrast2 = 'T2'): +def CreateDatasetSynthesis(phase, input_path, contrast1 = 'T1', contrast2 = 'T2', + image_size = 256): target_file = input_path + "/data_{}_{}.mat".format(phase, contrast1) - data_fs_s1=LoadDataSet(target_file) + data_fs_s1=LoadDataSet(target_file, target_size=image_size) target_file = input_path + "/data_{}_{}.mat".format(phase, contrast2) - data_fs_s2=LoadDataSet(target_file) + data_fs_s2=LoadDataSet(target_file, target_size=image_size) if data_fs_s1.shape[0] != data_fs_s2.shape[0]: raise ValueError( @@ -24,7 +25,7 @@ def CreateDatasetSynthesis(phase, input_path, contrast1 = 'T1', contrast2 = 'T2' -#Dataset loading from load_dir and converintg to 256x256 +#Dataset loading from load_dir, zero-padded out to target_size squared def LoadDataSet(load_dir, variable = 'data_fs', padding = True, Norm = True, target_size = 256): if not os.path.isfile(load_dir): diff --git a/test.py b/test.py index ee58641c..a873b8f4 100644 --- a/test.py +++ b/test.py @@ -54,7 +54,7 @@ def sample_and_test(args): #loading dataset phase='test' - dataset=CreateDatasetSynthesis('test', args.input_path, args.contrast1, args.contrast2) + dataset=CreateDatasetSynthesis('test', args.input_path, args.contrast1, args.contrast2, image_size=args.image_size) data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, diff --git a/tests/test_dataset.py b/tests/test_dataset.py index d1f2e366..640d79a3 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -86,3 +86,26 @@ def test_mismatched_slice_counts_are_rejected(tmp_path): write_mat(tmp_path / 'data_train_T2.mat', np.zeros((3, 152, 256), np.float32)) with pytest.raises(ValueError, match="different number of slices"): CreateDatasetSynthesis('train', str(tmp_path), 'T1', 'T2') + + +@pytest.mark.parametrize("image_size", [128, 192, 256]) +def test_target_size_is_configurable(tmp_path, image_size): + """--image_size has to reach the loader, not just the network.""" + raw = np.zeros((2, 100, 120), dtype=np.float32) + data = LoadDataSet(write_mat(tmp_path / 'a.mat', raw), target_size=image_size) + assert data.shape == (2, 1, image_size, image_size) + + +def test_create_dataset_forwards_the_image_size(tmp_path): + for contrast in ('T1', 'T2'): + write_mat(tmp_path / 'data_train_{}.mat'.format(contrast), + np.zeros((2, 100, 120), np.float32)) + dataset = CreateDatasetSynthesis('train', str(tmp_path), 'T1', 'T2', image_size=128) + first, second = dataset[0] + assert first.shape == second.shape == (1, 128, 128) + + +def test_input_larger_than_the_target_is_rejected(tmp_path): + raw = np.zeros((2, 200, 256), dtype=np.float32) + with pytest.raises(ValueError, match="does not fit"): + LoadDataSet(write_mat(tmp_path / 'a.mat', raw), target_size=128) diff --git a/tests/test_models.py b/tests/test_models.py index 45d99c78..fcab839f 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -130,3 +130,18 @@ def test_translation_networks_preserve_spatial_shape(): disc = generator_resnet.define_D(gpu_ids=[]) # PatchGAN returns a map of patch scores rather than one scalar assert disc(out).ndim == 4 + + +@pytest.mark.parametrize("ngf", [16, 32, 64]) +def test_define_g_honours_the_channel_width(ngf): + """--ngf must reach the translation networks, not just the discriminators.""" + net = generator_resnet.define_G(netG='resnet_6blocks', ngf=ngf, gpu_ids=[]) + # the first conv after the reflection pad produces ngf feature maps + first_conv = [m for m in net.modules() if isinstance(m, nn.Conv2d)][0] + assert first_conv.out_channels == ngf + + +def test_define_d_honours_the_channel_width(): + net = generator_resnet.define_D(ndf=16, gpu_ids=[]) + first_conv = [m for m in net.modules() if isinstance(m, nn.Conv2d)][0] + assert first_conv.out_channels == 16 diff --git a/train.py b/train.py index 0e53ce77..c8463612 100644 --- a/train.py +++ b/train.py @@ -104,8 +104,8 @@ def train_syndiff(rank, gpu, args): nz = args.nz #latent dimension - dataset = CreateDatasetSynthesis(phase = "train", input_path = args.input_path, contrast1 = args.contrast1, contrast2 = args.contrast2) - dataset_val = CreateDatasetSynthesis(phase = "val", input_path = args.input_path, contrast1 = args.contrast1, contrast2 = args.contrast2 ) + dataset = CreateDatasetSynthesis(phase = "train", input_path = args.input_path, contrast1 = args.contrast1, contrast2 = args.contrast2, image_size = args.image_size) + dataset_val = CreateDatasetSynthesis(phase = "val", input_path = args.input_path, contrast1 = args.contrast1, contrast2 = args.contrast2, image_size = args.image_size ) @@ -142,8 +142,8 @@ def train_syndiff(rank, gpu, args): gen_diffusive_1 = NCSNpp(args).to(device) gen_diffusive_2 = NCSNpp(args).to(device) #networks performing translation - gen_non_diffusive_1to2 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=gpu_ids) - gen_non_diffusive_2to1 = backbones.generator_resnet.define_G(netG='resnet_6blocks',gpu_ids=gpu_ids) + gen_non_diffusive_1to2 = backbones.generator_resnet.define_G(netG='resnet_6blocks',ngf=args.ngf,gpu_ids=gpu_ids) + gen_non_diffusive_2to1 = backbones.generator_resnet.define_G(netG='resnet_6blocks',ngf=args.ngf,gpu_ids=gpu_ids) disc_diffusive_1 = Discriminator_large(nc = 2, ngf = args.ngf, t_emb_dim = args.t_emb_dim, @@ -152,8 +152,8 @@ def train_syndiff(rank, gpu, args): t_emb_dim = args.t_emb_dim, act=nn.LeakyReLU(0.2)).to(device) - disc_non_diffusive_cycle1 = backbones.generator_resnet.define_D(gpu_ids=gpu_ids) - disc_non_diffusive_cycle2 = backbones.generator_resnet.define_D(gpu_ids=gpu_ids) + disc_non_diffusive_cycle1 = backbones.generator_resnet.define_D(ndf=args.ngf,gpu_ids=gpu_ids) + disc_non_diffusive_cycle2 = backbones.generator_resnet.define_D(ndf=args.ngf,gpu_ids=gpu_ids) broadcast_params(gen_diffusive_1.parameters()) broadcast_params(gen_diffusive_2.parameters()) From 37f23c0912c279a21b8943c31a5a891327d54452 Mon Sep 17 00:00:00 2001 From: Mertcan Ozdemir Date: Mon, 17 Aug 2026 13:19:40 +0300 Subject: [PATCH 21/21] docs: describe what --ngf and --image_size now cover Both arguments reach the data loader and the translation networks as of the previous commit; say so, and note that the documented values match the constants they replaced. --- README.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 02160b32..cb7c8d01 100644 --- a/README.md +++ b/README.md @@ -66,8 +66,8 @@ to `--contrast1` and `--contrast2`. Each `.mat` file is an HDF5 file holding a single variable named `data_fs` of shape `(#images, width, height)`, with image values in roughly `[0, 1]`. -Volumes are zero-padded out to 256x256 on load and rescaled to `[-1, 1]`, so -neither dimension may exceed 256. +Volumes are zero-padded out to `--image_size` squared on load and rescaled to +`[-1, 1]`, so neither dimension may exceed that size. ### Sample Data Sample toy data can be found under the `SynDiff_sample_data` folder. Note that @@ -85,6 +85,11 @@ val / test parts and name the parts as shown above. python3 train.py --image_size 256 --exp exp_syndiff --num_channels 2 --num_channels_dae 64 --ch_mult 1 1 2 2 4 4 --num_timesteps 4 --num_res_blocks 2 --batch_size 1 --contrast1 T1 --contrast2 T2 --num_epoch 500 --ngf 64 --embedding_type positional --use_ema --ema_decay 0.999 --r1_gamma 1. --z_emb_dim 256 --lr_d 1e-4 --lr_g 1.6e-4 --lazy_reg 10 --num_process_per_node 1 --save_content --local_rank 0 --input_path /input/path/for/data --output_path /output/for/results ``` +`--ngf` sets the channel width of the discriminators *and* of the translation +networks, and `--image_size` sets both the network resolution and the grid the +input volumes are padded to. The values above, `64` and `256`, are the ones +used in the paper. + `--num_process_per_node` controls the number of processes. With more than one a NCCL process group is set up and the networks are wrapped in `DistributedDataParallel`; with a single process neither is used, and the run