diff --git a/.DS_Store b/.DS_Store deleted file mode 100644 index f51ca32b..00000000 Binary files a/.DS_Store and /dev/null differ 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/README.md b/README.md index f271c74e..cb7c8d01 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 `--image_size` squared on load and rescaled to +`[-1, 1]`, so neither dimension may exceed that size. + ### 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,16 @@ 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 ``` +`--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 +falls back to CPU when no GPU is visible. +
## Pretrained Models @@ -70,6 +108,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. +

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 ) 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/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/dataset.py b/dataset.py index 46a75694..ce8f63b7 100644 --- a/dataset.py +++ b/dataset.py @@ -1,34 +1,69 @@ +import os + import torch.utils.data import numpy as np, h5py 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( + "'{}' 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=torch.utils.data.TensorDataset(torch.from_numpy(data_fs_s1),torch.from_numpy(data_fs_s2)) - return dataset +#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): + 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 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/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 diff --git a/test.py b/test.py index e79bcd47..a873b8f4 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,144 +21,40 @@ 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) checkpoint = torch.load(checkpoint_file, map_location=device) ckpt = checkpoint - - for key in list(ckpt.keys()): - ckpt[key[7:]] = ckpt.pop(key) - netG.load_state_dict(ckpt) + + # 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.' + 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): 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. #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, @@ -180,13 +78,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) @@ -204,7 +107,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)) @@ -230,7 +133,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)) @@ -243,6 +146,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) @@ -332,6 +238,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) 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..640d79a3 --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,111 @@ +"""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') + + +@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_diffusion.py b/tests/test_diffusion.py new file mode 100644 index 00000000..1fcb3fba --- /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 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): + 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..fcab839f --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,147 @@ +"""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 + + +@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/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) diff --git a/train.py b/train.py index e2629dea..c8463612 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 @@ -26,160 +29,52 @@ 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) -#%% 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 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 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) +def maybe_ddp(model, device_ids): + """Wrap in DistributedDataParallel only when a process group is active. - return out + 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 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 strip_module_prefix(state_dict): + """Drop the 'module.' prefixes that parallel wrappers add to every key. -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} + 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. """ - 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 + prefix = 'module.' + 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)) + + #%% def train_syndiff(rank, gpu, args): @@ -197,17 +92,20 @@ 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 - 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 ) @@ -232,8 +130,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. @@ -242,9 +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 - 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',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, @@ -253,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]) - disc_non_diffusive_cycle2 = backbones.generator_resnet.define_D(gpu_ids=[gpu]) + 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()) @@ -299,15 +198,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 @@ -326,13 +225,15 @@ 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']) - 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']) @@ -344,19 +245,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'] @@ -447,14 +348,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) @@ -489,8 +393,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) @@ -577,10 +482,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() @@ -696,9 +599,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() @@ -718,15 +621,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') 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: 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