├── IRM ├── improved_diffusion │ ├── __init__.py │ ├── dist_util.py │ ├── fp16_util.py │ ├── gaussian_diffusion.py │ ├── helper.py │ ├── image_datasets.py │ ├── logger.py │ ├── losses.py │ ├── nn.py │ ├── reco_train.py │ ├── resample.py │ ├── respace.py │ ├── script_util.py │ ├── scripts │ │ ├── image_nll.py │ │ ├── img_multi_sample.py │ │ ├── img_super_res_sample.py │ │ └── img_super_res_train.py │ ├── train_util.py │ └── unet.py └── setup.py ├── README.md └── SUM ├── __init__.py ├── dist_util.py ├── fp16_util.py ├── gaussian_diffusion.py ├── helper.py ├── image_datasets.py ├── logger.py ├── losses.py ├── nn.py ├── reco_train.py ├── resample.py ├── respace.py ├── script_util.py ├── scripts ├── image_nll.py ├── multi_super_res.py ├── super_res_sample.py └── super_res_train.py ├── train_util.py └── unet.py /IRM/improved_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = 8 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | 28 | comm = MPI.COMM_WORLD 29 | backend = "gloo" if not th.cuda.is_available() else "nccl" 30 | 31 | if backend == "gloo": 32 | hostname = "localhost" 33 | else: 34 | hostname = socket.gethostbyname(socket.getfqdn()) 35 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 36 | os.environ["RANK"] = str(comm.rank) 37 | os.environ["WORLD_SIZE"] = str(comm.size) 38 | 39 | port = comm.bcast(_find_free_port(), root=0) 40 | os.environ["MASTER_PORT"] = str(port) 41 | dist.init_process_group(backend=backend, init_method="env://") 42 | 43 | 44 | def dev(): 45 | """ 46 | Get the device to use for torch.distributed. 47 | """ 48 | if th.cuda.is_available(): 49 | return th.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}") 50 | return th.device("cpu") 51 | 52 | 53 | def load_state_dict(path, **kwargs): 54 | """ 55 | Load a PyTorch file without redundant fetches across MPI ranks. 56 | """ 57 | if MPI.COMM_WORLD.Get_rank() == 0: 58 | with bf.BlobFile(path, "rb") as f: 59 | data = f.read() 60 | else: 61 | data = None 62 | data = MPI.COMM_WORLD.bcast(data) 63 | return th.load(io.BytesIO(data), **kwargs) 64 | 65 | 66 | def sync_params(params): 67 | """ 68 | Synchronize a sequence of Tensors across ranks from rank 0. 69 | """ 70 | for p in params: 71 | with th.no_grad(): 72 | dist.broadcast(p, 0) 73 | 74 | 75 | def _find_free_port(): 76 | try: 77 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 78 | s.bind(("", 0)) 79 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 80 | return s.getsockname()[1] 81 | finally: 82 | s.close() 83 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import torch.nn as nn 6 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 7 | 8 | 9 | def convert_module_to_f16(l): 10 | """ 11 | Convert primitive modules to float16. 12 | """ 13 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 14 | l.weight.data = l.weight.data.half() 15 | l.bias.data = l.bias.data.half() 16 | 17 | 18 | def convert_module_to_f32(l): 19 | """ 20 | Convert primitive modules to float32, undoing convert_module_to_f16(). 21 | """ 22 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 23 | l.weight.data = l.weight.data.float() 24 | l.bias.data = l.bias.data.float() 25 | 26 | 27 | def make_master_params(model_params): 28 | """ 29 | Copy model parameters into a (differently-shaped) list of full-precision 30 | parameters. 31 | """ 32 | master_params = _flatten_dense_tensors( 33 | [param.detach().float() for param in model_params] 34 | ) 35 | master_params = nn.Parameter(master_params) 36 | master_params.requires_grad = True 37 | return [master_params] 38 | 39 | 40 | def model_grads_to_master_grads(model_params, master_params): 41 | """ 42 | Copy the gradients from the model parameters into the master parameters 43 | from make_master_params(). 44 | """ 45 | master_params[0].grad = _flatten_dense_tensors( 46 | [param.grad.data.detach().float() for param in model_params] 47 | ) 48 | 49 | 50 | def master_params_to_model_params(model_params, master_params): 51 | """ 52 | Copy the master parameter data back into the model parameters. 53 | """ 54 | # Without copying to a list, if a generator is passed, this will 55 | # silently not copy any parameters. 56 | model_params = list(model_params) 57 | 58 | for param, master_param in zip( 59 | model_params, unflatten_master_params(model_params, master_params) 60 | ): 61 | param.detach().copy_(master_param) 62 | 63 | 64 | def unflatten_master_params(model_params, master_params): 65 | """ 66 | Unflatten the master parameters to look like model_params. 67 | """ 68 | return _unflatten_dense_tensors(master_params[0].detach(), model_params) 69 | 70 | 71 | def zero_grad(model_params): 72 | for param in model_params: 73 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 74 | if param.grad is not None: 75 | param.grad.detach_() 76 | param.grad.zero_() 77 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/helper.py: -------------------------------------------------------------------------------- 1 | from skimage.io import imread, imsave 2 | import numpy as np 3 | import pandas as pd 4 | import json 5 | import tifffile 6 | 7 | 8 | def load_tiff_stack(file): 9 | ''' 10 | 11 | :param file: Path object describing the location of the file 12 | :return: a numpy array of the volume 13 | ''' 14 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 15 | raise FileNotFoundError('Input has to be tif.') 16 | else: 17 | volume = imread(file, plugin='tifffile') 18 | return volume 19 | 20 | 21 | def load_tiff_stack_with_metadata(file): 22 | ''' 23 | 24 | :param file: Path object describing the location of the file 25 | :return: a numpy array of the volume, a dict with the metadata 26 | ''' 27 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 28 | raise FileNotFoundError('File has to be tif.') 29 | with tifffile.TiffFile(file) as tif: 30 | data = tif.asarray() 31 | metadata = tif.pages[0].tags["ImageDescription"].value 32 | metadata = metadata.replace("'", "\"") 33 | try: 34 | metadata = json.loads(metadata) 35 | except: 36 | print('The tiff file you try to open does not seem to have metadata attached.') 37 | metadata = None 38 | return data, metadata 39 | 40 | 41 | def save_to_tiff_stack(array, file): 42 | ''' 43 | 44 | :param array: Array to save 45 | :param file: Path object to save to 46 | ''' 47 | if not file.parent.is_dir(): 48 | file.parent.mkdir(parents=True, exist_ok=True) 49 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 50 | raise FileNotFoundError('File has to be tif.') 51 | else: 52 | imsave(file, array, plugin='tifffile', check_contrast=False) 53 | 54 | 55 | def save_to_tiff_stack_with_metadata(array, file, metadata): 56 | ''' 57 | 58 | :param array: 59 | :param file: 60 | :return: 61 | ''' 62 | if not file.parent.is_dir(): 63 | file.parent.mkdir(parents=True, exist_ok=True) 64 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 65 | raise FileNotFoundError('File has to be tif.') 66 | else: 67 | # metadata = json.dumps(metadata) 68 | # tifffile.imsave(file, array, description=metadata) 69 | tifffile.imwrite(file, shape=array.shape, dtype=array.dtype, metadata=metadata) 70 | 71 | # memory map numpy array to data in OME-TIFF file 72 | memmap_stack = tifffile.memmap(file) 73 | 74 | # write data to memory-mapped array 75 | for t in range(array.shape[0]): 76 | memmap_stack[t] = array[t, :, :] 77 | memmap_stack.flush() 78 | 79 | 80 | def load_csv(file): 81 | ''' 82 | 83 | :param file: Path object describing the location of the file 84 | :return: pandas dataframe containing the information from the file 85 | ''' 86 | if not file.name.endswith('.csv'): 87 | raise FileNotFoundError('Input has to be a csv file.') 88 | else: 89 | data = pd.read_csv(file) 90 | return data 91 | 92 | 93 | def save_to_json(params, file): 94 | ''' 95 | 96 | :param params: parameter dict 97 | :param file: path to a file to save to 98 | :return: 99 | ''' 100 | if not file.parent.is_dir(): 101 | file.parent.mkdir(parents=True, exist_ok=True) 102 | with open(file, 'w') as f: 103 | json.dump(params, f, indent=2) 104 | 105 | 106 | def load_from_json(file): 107 | ''' 108 | 109 | :param file: path to a file to load 110 | :return: contents of the file 111 | ''' 112 | with open(file) as f: 113 | data = json.load(f) 114 | return data 115 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/image_datasets.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import blobfile as bf 3 | from mpi4py import MPI 4 | import numpy as np 5 | from torch.utils.data import DataLoader, Dataset 6 | import torch 7 | import os 8 | import random 9 | import torch.nn.functional as F 10 | 11 | 12 | 13 | def improved_data_preprocess(sino_path = '/root/autodl-tmp/TASK4/SINO_DATA/val_stage1', batch_size = 1, shuffle = True, num_workers = 8): 14 | train_dataset = diffCT_Dataset(sino_path) 15 | train_data_loader = DataLoader(train_dataset, batch_size, shuffle, num_workers=num_workers, pin_memory=True, drop_last=True) 16 | while True: 17 | yield from train_data_loader 18 | 19 | def count_npy_files(folder_path: str): 20 | files = os.listdir(folder_path) 21 | npy_files = [file for file in files if file.endswith('.npy')] 22 | return len(npy_files) 23 | 24 | class diffCT_Dataset(Dataset): 25 | def __init__(self, sino_path): 26 | self.sino_path = sino_path 27 | self.sino_num = count_npy_files(sino_path) 28 | def __len__(self): 29 | return 3240 + 4140 +3260 + 3260 #self.sino_num # C + L 3254 + 4152 30 | def __getitem__(self,index): 31 | out_dict = {} 32 | if index<3240: 33 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load(self.sino_path + f'/full_C_parallel_img/{index}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 34 | # x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_C_parallel_sino/{index}.npy')).to(torch.float32),dim=0) 35 | out_dict["low_res"] = torch.tensor(np.load(self.sino_path + f'/36to288_C_parallel_img/{index}.npy')).to(torch.float32) 36 | out_dict["_36_res"] = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_C_parallel_sino/{index}.npy')).to(torch.float32),dim=0) 37 | # out_dict["y"] = np.array(0, dtype=np.int64) 38 | 39 | elif index<(3240 + 4140): 40 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load(self.sino_path + f'/full_L_parallel_img/{index - 3240}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 41 | # x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index - 3254}.npy')).to(torch.float32),dim=0) 42 | out_dict["low_res"] = torch.tensor(np.load(self.sino_path + f'/36to288_L_parallel_img/{index - 3240}.npy')).to(torch.float32) 43 | out_dict["_36_res"] = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index - 3240}.npy')).to(torch.float32),dim=0) 44 | # out_dict["y"] = np.array(1, dtype=np.int64) 45 | 46 | elif index < (3240 + 4140 + 3260): 47 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/full_C_parallel_img/{index - 3240 - 4140}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 48 | out_dict["low_res"] = torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/36to288_C_parallel_img/{index - 3240 - 4140}.npy')).to(torch.float32) 49 | out_dict["_36_res"] = torch.unsqueeze(torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/full_C_parallel_sino/{index - 3240 - 4140}.npy')).to(torch.float32),dim=0) 50 | 51 | else: 52 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/full_L_parallel_img/{index - 3240 - 4140 - 3260}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 53 | out_dict["low_res"] = torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/36to288_L_parallel_img/{index - 3240 - 4140 - 3260}.npy')).to(torch.float32) 54 | out_dict["_36_res"] = torch.unsqueeze(torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/full_L_parallel_sino/{index - 3240 - 4140 - 3260}.npy')).to(torch.float32),dim=0) 55 | 56 | # x_nd_mean = torch.mean(x_nd) 57 | # x_nd = x_nd - x_nd_mean 58 | # x_nd_min_abs = torch.abs(torch.min(x_nd)) 59 | # x_nd = x_nd / x_nd_min_abs 60 | 61 | # out_dict["low_res"] = out_dict["low_res"] - x_nd_mean # torch.mean(out_dict["low_res"]) 62 | # out_dict["low_res"] = out_dict["low_res"] / x_nd_min_abs # torch.abs(torch.min(out_dict["low_res"])) 63 | # print("x_nd:", torch.min(x_nd)) 64 | # print("out_dict:", torch.min(out_dict["low_res"])) 65 | # print("out_dict:", torch.mean(out_dict["low_res"])) 66 | 67 | assert len(x_nd.shape) == 3 68 | return x_nd, out_dict 69 | 70 | 71 | 72 | 73 | 74 | def improved_data_preprocess_val(sino_path = '/root/autodl-tmp/TASK4/SINO_DATA/val_stage2', batch_size = 1, shuffle = False, num_workers = 8): 75 | train_dataset = diffCT_Dataset_val(sino_path) 76 | train_data_loader = DataLoader(train_dataset, batch_size, shuffle, num_workers=num_workers, pin_memory=True, drop_last=True) 77 | while True: 78 | yield from train_data_loader 79 | 80 | 81 | class diffCT_Dataset_val(Dataset): 82 | def __init__(self, sino_path): 83 | self.sino_path = sino_path 84 | self.sino_num = count_npy_files(sino_path) 85 | def __len__(self): 86 | return 3240 + 4140 +3260 + 3260 #self.sino_num # C + L 87 | def __getitem__(self,index): 88 | out_dict = {} 89 | index = 150 90 | if index<3240: 91 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load(self.sino_path + f'/full_C_parallel_img/{index}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 92 | # x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_C_parallel_sino/{index}.npy')).to(torch.float32),dim=0) 93 | out_dict["low_res"] = torch.tensor(np.load(self.sino_path + f'/36to288_C_parallel_img/{index}.npy')).to(torch.float32) 94 | # out_dict["_72_res"] = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_C_parallel_sino/{index}.npy')).to(torch.float32),dim=0) 95 | # out_dict["y"] = np.array(0, dtype=np.int64) 96 | 97 | elif index<(3240 + 4140): 98 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load(self.sino_path + f'/full_L_parallel_img/{index - 3240}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 99 | # x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index - 3254}.npy')).to(torch.float32),dim=0) 100 | out_dict["low_res"] = torch.tensor(np.load(self.sino_path + f'/36to288_L_parallel_img/{index - 3240}.npy')).to(torch.float32) 101 | # out_dict["_72_res"] = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index - 3254}.npy')).to(torch.float32),dim=0) 102 | # out_dict["y"] = np.array(1, dtype=np.int64) 103 | 104 | elif index < (3240 + 4140 + 3260): 105 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/full_C_parallel_img/{index - 3240 - 4140}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 106 | out_dict["low_res"] = torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/36to288_C_parallel_img/{index - 3240 - 4140}.npy')).to(torch.float32) 107 | 108 | else: 109 | x_nd = torch.flip(torch.unsqueeze(100. * torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/full_L_parallel_img/{index - 3240 - 4140 - 3260}.npy')[112:624,112:624]).to(torch.float32),dim=0),dims=[1,2]) 110 | out_dict["low_res"] = torch.tensor(np.load('/root/autodl-tmp/TASK4/SINO_DATA/train_stage1' + f'/36to288_L_parallel_img/{index - 3240 - 4140 - 3260}.npy')).to(torch.float32) 111 | 112 | assert len(x_nd.shape) == 3 113 | return x_nd, out_dict 114 | 115 | 116 | 117 | 118 | def trans_data_preprocess(sino_path = '/root/autodl-tmp/TASK4/SINO_DATA/val_stage2', batch_size = 1, shuffle = False, num_workers = 8): 119 | train_dataset = trans_diffCT_Dataset(sino_path) 120 | train_data_loader = DataLoader(train_dataset, batch_size, shuffle, num_workers=num_workers, pin_memory=True, drop_last=True) 121 | while True: 122 | yield from train_data_loader 123 | 124 | 125 | class trans_diffCT_Dataset(Dataset): 126 | def __init__(self, sino_path): 127 | self.sino_path = sino_path 128 | self.sino_num = count_npy_files(sino_path) 129 | def __len__(self): 130 | # return 2020-496 131 | return 800# -496 132 | def __getitem__(self,index): 133 | out_dict = {} 134 | index = index +0# +496 +450/98 +103 135 | x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index}.npy')).to(torch.float32),dim=0) 136 | out_dict["low_res"] = torch.tensor(np.load(self.sino_path + f'/test_pic_72/{index}.npy')).to(torch.float32) # 36to288_C_parallel_img test_pic # 36to288_L_parallel_img 137 | out_dict["_36_res"] = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index}.npy')).to(torch.float32),dim=0) 138 | 139 | # out_dict["y"] = np.array(1, dtype=np.int64) 140 | 141 | out_dict["index"] = index 142 | assert len(x_nd.shape) == 3 143 | return x_nd, out_dict -------------------------------------------------------------------------------- /IRM/improved_diffusion/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Logger copied from OpenAI baselines to avoid extra RL-based dependencies: 3 | https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py 4 | """ 5 | 6 | import os 7 | import sys 8 | import shutil 9 | import os.path as osp 10 | import json 11 | import time 12 | import datetime 13 | import tempfile 14 | import warnings 15 | from collections import defaultdict 16 | from contextlib import contextmanager 17 | 18 | DEBUG = 10 19 | INFO = 20 20 | WARN = 30 21 | ERROR = 40 22 | 23 | DISABLED = 50 24 | 25 | 26 | class KVWriter(object): 27 | def writekvs(self, kvs): 28 | raise NotImplementedError 29 | 30 | 31 | class SeqWriter(object): 32 | def writeseq(self, seq): 33 | raise NotImplementedError 34 | 35 | 36 | class HumanOutputFormat(KVWriter, SeqWriter): 37 | def __init__(self, filename_or_file): 38 | if isinstance(filename_or_file, str): 39 | self.file = open(filename_or_file, "wt") 40 | self.own_file = True 41 | else: 42 | assert hasattr(filename_or_file, "read"), ( 43 | "expected file or str, got %s" % filename_or_file 44 | ) 45 | self.file = filename_or_file 46 | self.own_file = False 47 | 48 | def writekvs(self, kvs): 49 | # Create strings for printing 50 | key2str = {} 51 | for (key, val) in sorted(kvs.items()): 52 | if hasattr(val, "__float__"): 53 | valstr = "%-8.3g" % val 54 | else: 55 | valstr = str(val) 56 | key2str[self._truncate(key)] = self._truncate(valstr) 57 | 58 | # Find max widths 59 | if len(key2str) == 0: 60 | print("WARNING: tried to write empty key-value dict") 61 | return 62 | else: 63 | keywidth = max(map(len, key2str.keys())) 64 | valwidth = max(map(len, key2str.values())) 65 | 66 | # Write out the data 67 | dashes = "-" * (keywidth + valwidth + 7) 68 | lines = [dashes] 69 | for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): 70 | lines.append( 71 | "| %s%s | %s%s |" 72 | % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) 73 | ) 74 | lines.append(dashes) 75 | self.file.write("\n".join(lines) + "\n") 76 | 77 | # Flush the output to the file 78 | self.file.flush() 79 | 80 | def _truncate(self, s): 81 | maxlen = 30 82 | return s[: maxlen - 3] + "..." if len(s) > maxlen else s 83 | 84 | def writeseq(self, seq): 85 | seq = list(seq) 86 | for (i, elem) in enumerate(seq): 87 | self.file.write(elem) 88 | if i < len(seq) - 1: # add space unless this is the last one 89 | self.file.write(" ") 90 | self.file.write("\n") 91 | self.file.flush() 92 | 93 | def close(self): 94 | if self.own_file: 95 | self.file.close() 96 | 97 | 98 | class JSONOutputFormat(KVWriter): 99 | def __init__(self, filename): 100 | self.file = open(filename, "wt") 101 | 102 | def writekvs(self, kvs): 103 | for k, v in sorted(kvs.items()): 104 | if hasattr(v, "dtype"): 105 | kvs[k] = float(v) 106 | self.file.write(json.dumps(kvs) + "\n") 107 | self.file.flush() 108 | 109 | def close(self): 110 | self.file.close() 111 | 112 | 113 | class CSVOutputFormat(KVWriter): 114 | def __init__(self, filename): 115 | self.file = open(filename, "w+t") 116 | self.keys = [] 117 | self.sep = "," 118 | 119 | def writekvs(self, kvs): 120 | # Add our current row to the history 121 | extra_keys = list(kvs.keys() - self.keys) 122 | extra_keys.sort() 123 | if extra_keys: 124 | self.keys.extend(extra_keys) 125 | self.file.seek(0) 126 | lines = self.file.readlines() 127 | self.file.seek(0) 128 | for (i, k) in enumerate(self.keys): 129 | if i > 0: 130 | self.file.write(",") 131 | self.file.write(k) 132 | self.file.write("\n") 133 | for line in lines[1:]: 134 | self.file.write(line[:-1]) 135 | self.file.write(self.sep * len(extra_keys)) 136 | self.file.write("\n") 137 | for (i, k) in enumerate(self.keys): 138 | if i > 0: 139 | self.file.write(",") 140 | v = kvs.get(k) 141 | if v is not None: 142 | self.file.write(str(v)) 143 | self.file.write("\n") 144 | self.file.flush() 145 | 146 | def close(self): 147 | self.file.close() 148 | 149 | 150 | class TensorBoardOutputFormat(KVWriter): 151 | """ 152 | Dumps key/value pairs into TensorBoard's numeric format. 153 | """ 154 | 155 | def __init__(self, dir): 156 | os.makedirs(dir, exist_ok=True) 157 | self.dir = dir 158 | self.step = 1 159 | prefix = "events" 160 | path = osp.join(osp.abspath(dir), prefix) 161 | import tensorflow as tf 162 | from tensorflow.python import pywrap_tensorflow 163 | from tensorflow.core.util import event_pb2 164 | from tensorflow.python.util import compat 165 | 166 | self.tf = tf 167 | self.event_pb2 = event_pb2 168 | self.pywrap_tensorflow = pywrap_tensorflow 169 | self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) 170 | 171 | def writekvs(self, kvs): 172 | def summary_val(k, v): 173 | kwargs = {"tag": k, "simple_value": float(v)} 174 | return self.tf.Summary.Value(**kwargs) 175 | 176 | summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) 177 | event = self.event_pb2.Event(wall_time=time.time(), summary=summary) 178 | event.step = ( 179 | self.step 180 | ) # is there any reason why you'd want to specify the step? 181 | self.writer.WriteEvent(event) 182 | self.writer.Flush() 183 | self.step += 1 184 | 185 | def close(self): 186 | if self.writer: 187 | self.writer.Close() 188 | self.writer = None 189 | 190 | 191 | def make_output_format(format, ev_dir, log_suffix=""): 192 | os.makedirs(ev_dir, exist_ok=True) 193 | if format == "stdout": 194 | return HumanOutputFormat(sys.stdout) 195 | elif format == "log": 196 | return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) 197 | elif format == "json": 198 | return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) 199 | elif format == "csv": 200 | return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) 201 | elif format == "tensorboard": 202 | return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) 203 | else: 204 | raise ValueError("Unknown format specified: %s" % (format,)) 205 | 206 | 207 | # ================================================================ 208 | # API 209 | # ================================================================ 210 | 211 | 212 | def logkv(key, val): 213 | """ 214 | Log a value of some diagnostic 215 | Call this once for each diagnostic quantity, each iteration 216 | If called many times, last value will be used. 217 | """ 218 | get_current().logkv(key, val) 219 | 220 | 221 | def logkv_mean(key, val): 222 | """ 223 | The same as logkv(), but if called many times, values averaged. 224 | """ 225 | get_current().logkv_mean(key, val) 226 | 227 | 228 | def logkvs(d): 229 | """ 230 | Log a dictionary of key-value pairs 231 | """ 232 | for (k, v) in d.items(): 233 | logkv(k, v) 234 | 235 | 236 | def dumpkvs(): 237 | """ 238 | Write all of the diagnostics from the current iteration 239 | """ 240 | return get_current().dumpkvs() 241 | 242 | 243 | def getkvs(): 244 | return get_current().name2val 245 | 246 | 247 | def log(*args, level=INFO): 248 | """ 249 | Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). 250 | """ 251 | get_current().log(*args, level=level) 252 | 253 | 254 | def debug(*args): 255 | log(*args, level=DEBUG) 256 | 257 | 258 | def info(*args): 259 | log(*args, level=INFO) 260 | 261 | 262 | def warn(*args): 263 | log(*args, level=WARN) 264 | 265 | 266 | def error(*args): 267 | log(*args, level=ERROR) 268 | 269 | 270 | def set_level(level): 271 | """ 272 | Set logging threshold on current logger. 273 | """ 274 | get_current().set_level(level) 275 | 276 | 277 | def set_comm(comm): 278 | get_current().set_comm(comm) 279 | 280 | 281 | def get_dir(): 282 | """ 283 | Get directory that log files are being written to. 284 | will be None if there is no output directory (i.e., if you didn't call start) 285 | """ 286 | return get_current().get_dir() 287 | 288 | 289 | record_tabular = logkv 290 | dump_tabular = dumpkvs 291 | 292 | 293 | @contextmanager 294 | def profile_kv(scopename): 295 | logkey = "wait_" + scopename 296 | tstart = time.time() 297 | try: 298 | yield 299 | finally: 300 | get_current().name2val[logkey] += time.time() - tstart 301 | 302 | 303 | def profile(n): 304 | """ 305 | Usage: 306 | @profile("my_func") 307 | def my_func(): code 308 | """ 309 | 310 | def decorator_with_name(func): 311 | def func_wrapper(*args, **kwargs): 312 | with profile_kv(n): 313 | return func(*args, **kwargs) 314 | 315 | return func_wrapper 316 | 317 | return decorator_with_name 318 | 319 | 320 | # ================================================================ 321 | # Backend 322 | # ================================================================ 323 | 324 | 325 | def get_current(): 326 | if Logger.CURRENT is None: 327 | _configure_default_logger() 328 | 329 | return Logger.CURRENT 330 | 331 | 332 | class Logger(object): 333 | DEFAULT = None # A logger with no output files. (See right below class definition) 334 | # So that you can still log to the terminal without setting up any output files 335 | CURRENT = None # Current logger being used by the free functions above 336 | 337 | def __init__(self, dir, output_formats, comm=None): 338 | self.name2val = defaultdict(float) # values this iteration 339 | self.name2cnt = defaultdict(int) 340 | self.level = INFO 341 | self.dir = dir 342 | self.output_formats = output_formats 343 | self.comm = comm 344 | 345 | # Logging API, forwarded 346 | # ---------------------------------------- 347 | def logkv(self, key, val): 348 | self.name2val[key] = val 349 | 350 | def logkv_mean(self, key, val): 351 | oldval, cnt = self.name2val[key], self.name2cnt[key] 352 | self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) 353 | self.name2cnt[key] = cnt + 1 354 | 355 | def dumpkvs(self): 356 | if self.comm is None: 357 | d = self.name2val 358 | else: 359 | d = mpi_weighted_mean( 360 | self.comm, 361 | { 362 | name: (val, self.name2cnt.get(name, 1)) 363 | for (name, val) in self.name2val.items() 364 | }, 365 | ) 366 | if self.comm.rank != 0: 367 | d["dummy"] = 1 # so we don't get a warning about empty dict 368 | out = d.copy() # Return the dict for unit testing purposes 369 | for fmt in self.output_formats: 370 | if isinstance(fmt, KVWriter): 371 | fmt.writekvs(d) 372 | self.name2val.clear() 373 | self.name2cnt.clear() 374 | return out 375 | 376 | def log(self, *args, level=INFO): 377 | if self.level <= level: 378 | self._do_log(args) 379 | 380 | # Configuration 381 | # ---------------------------------------- 382 | def set_level(self, level): 383 | self.level = level 384 | 385 | def set_comm(self, comm): 386 | self.comm = comm 387 | 388 | def get_dir(self): 389 | return self.dir 390 | 391 | def close(self): 392 | for fmt in self.output_formats: 393 | fmt.close() 394 | 395 | # Misc 396 | # ---------------------------------------- 397 | def _do_log(self, args): 398 | for fmt in self.output_formats: 399 | if isinstance(fmt, SeqWriter): 400 | fmt.writeseq(map(str, args)) 401 | 402 | 403 | def get_rank_without_mpi_import(): 404 | # check environment variables here instead of importing mpi4py 405 | # to avoid calling MPI_Init() when this module is imported 406 | for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: 407 | if varname in os.environ: 408 | return int(os.environ[varname]) 409 | return 0 410 | 411 | 412 | def mpi_weighted_mean(comm, local_name2valcount): 413 | """ 414 | Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 415 | Perform a weighted average over dicts that are each on a different node 416 | Input: local_name2valcount: dict mapping key -> (value, count) 417 | Returns: key -> mean 418 | """ 419 | all_name2valcount = comm.gather(local_name2valcount) 420 | if comm.rank == 0: 421 | name2sum = defaultdict(float) 422 | name2count = defaultdict(float) 423 | for n2vc in all_name2valcount: 424 | for (name, (val, count)) in n2vc.items(): 425 | try: 426 | val = float(val) 427 | except ValueError: 428 | if comm.rank == 0: 429 | warnings.warn( 430 | "WARNING: tried to compute mean on non-float {}={}".format( 431 | name, val 432 | ) 433 | ) 434 | else: 435 | name2sum[name] += val * count 436 | name2count[name] += count 437 | return {name: name2sum[name] / name2count[name] for name in name2sum} 438 | else: 439 | return {} 440 | 441 | 442 | def configure(dir=None, format_strs=None, comm=None, log_suffix=""): 443 | """ 444 | If comm is provided, average all numerical stats across that comm 445 | """ 446 | if dir is None: 447 | dir = os.getenv("OPENAI_LOGDIR") 448 | if dir is None: 449 | dir = osp.join( 450 | tempfile.gettempdir(), 451 | datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 452 | ) 453 | assert isinstance(dir, str) 454 | dir = os.path.expanduser(dir) 455 | os.makedirs(os.path.expanduser(dir), exist_ok=True) 456 | 457 | rank = get_rank_without_mpi_import() 458 | if rank > 0: 459 | log_suffix = log_suffix + "-rank%03i" % rank 460 | 461 | if format_strs is None: 462 | if rank == 0: 463 | format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") 464 | else: 465 | format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") 466 | format_strs = filter(None, format_strs) 467 | output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] 468 | 469 | Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) 470 | if output_formats: 471 | log("Logging to %s" % dir) 472 | 473 | 474 | def _configure_default_logger(): 475 | configure() 476 | Logger.DEFAULT = Logger.CURRENT 477 | 478 | 479 | def reset(): 480 | if Logger.CURRENT is not Logger.DEFAULT: 481 | Logger.CURRENT.close() 482 | Logger.CURRENT = Logger.DEFAULT 483 | log("Reset logger") 484 | 485 | 486 | @contextmanager 487 | def scoped_configure(dir=None, format_strs=None, comm=None): 488 | prevlogger = Logger.CURRENT 489 | configure(dir=dir, format_strs=format_strs, comm=comm) 490 | try: 491 | yield 492 | finally: 493 | Logger.CURRENT.close() 494 | Logger.CURRENT = prevlogger 495 | 496 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/losses.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for various likelihood-based losses. These are ported from the original 3 | Ho et al. diffusion models codebase: 4 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 5 | """ 6 | 7 | import numpy as np 8 | 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/nn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various utilities for neural networks. 3 | """ 4 | 5 | import math 6 | 7 | import torch as th 8 | import torch.nn as nn 9 | 10 | 11 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 12 | class SiLU(nn.Module): 13 | def forward(self, x): 14 | return x * th.sigmoid(x) 15 | 16 | 17 | class GroupNorm32(nn.GroupNorm): 18 | def forward(self, x): 19 | return super().forward(x.float()).type(x.dtype) 20 | 21 | 22 | def conv_nd(dims, *args, **kwargs): 23 | """ 24 | Create a 1D, 2D, or 3D convolution module. 25 | """ 26 | if dims == 1: 27 | return nn.Conv1d(*args, **kwargs) 28 | elif dims == 2: 29 | return nn.Conv2d(*args, **kwargs) 30 | elif dims == 3: 31 | return nn.Conv3d(*args, **kwargs) 32 | raise ValueError(f"unsupported dimensions: {dims}") 33 | 34 | 35 | def linear(*args, **kwargs): 36 | """ 37 | Create a linear module. 38 | """ 39 | return nn.Linear(*args, **kwargs) 40 | 41 | 42 | def avg_pool_nd(dims, *args, **kwargs): 43 | """ 44 | Create a 1D, 2D, or 3D average pooling module. 45 | """ 46 | if dims == 1: 47 | return nn.AvgPool1d(*args, **kwargs) 48 | elif dims == 2: 49 | return nn.AvgPool2d(*args, **kwargs) 50 | elif dims == 3: 51 | return nn.AvgPool3d(*args, **kwargs) 52 | raise ValueError(f"unsupported dimensions: {dims}") 53 | 54 | 55 | def update_ema(target_params, source_params, rate=0.99): 56 | """ 57 | Update target parameters to be closer to those of source parameters using 58 | an exponential moving average. 59 | 60 | :param target_params: the target parameter sequence. 61 | :param source_params: the source parameter sequence. 62 | :param rate: the EMA rate (closer to 1 means slower). 63 | """ 64 | for targ, src in zip(target_params, source_params): 65 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 66 | 67 | 68 | def zero_module(module): 69 | """ 70 | Zero out the parameters of a module and return it. 71 | """ 72 | for p in module.parameters(): 73 | p.detach().zero_() 74 | return module 75 | 76 | 77 | def scale_module(module, scale): 78 | """ 79 | Scale the parameters of a module and return it. 80 | """ 81 | for p in module.parameters(): 82 | p.detach().mul_(scale) 83 | return module 84 | 85 | 86 | def mean_flat(tensor): 87 | """ 88 | Take the mean over all non-batch dimensions. 89 | """ 90 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 91 | 92 | 93 | def normalization(channels): 94 | """ 95 | Make a standard normalization layer. 96 | 97 | :param channels: number of input channels. 98 | :return: an nn.Module for normalization. 99 | """ 100 | return GroupNorm32(32, channels) 101 | 102 | 103 | def timestep_embedding(timesteps, dim, max_period=10000): 104 | """ 105 | Create sinusoidal timestep embeddings. 106 | 107 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 108 | These may be fractional. 109 | :param dim: the dimension of the output. 110 | :param max_period: controls the minimum frequency of the embeddings. 111 | :return: an [N x dim] Tensor of positional embeddings. 112 | """ 113 | half = dim // 2 114 | freqs = th.exp( 115 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 116 | ).to(device=timesteps.device) 117 | args = timesteps[:, None].float() * freqs[None] 118 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 119 | if dim % 2: 120 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 121 | return embedding 122 | 123 | 124 | def checkpoint(func, inputs, params, flag): 125 | """ 126 | Evaluate a function without caching intermediate activations, allowing for 127 | reduced memory at the expense of extra compute in the backward pass. 128 | 129 | :param func: the function to evaluate. 130 | :param inputs: the argument sequence to pass to `func`. 131 | :param params: a sequence of parameters `func` depends on but does not 132 | explicitly take as arguments. 133 | :param flag: if False, disable gradient checkpointing. 134 | """ 135 | if flag: 136 | args = tuple(inputs) + tuple(params) 137 | return CheckpointFunction.apply(func, len(inputs), *args) 138 | else: 139 | return func(*inputs) 140 | 141 | 142 | class CheckpointFunction(th.autograd.Function): 143 | @staticmethod 144 | def forward(ctx, run_function, length, *args): 145 | ctx.run_function = run_function 146 | ctx.input_tensors = list(args[:length]) 147 | ctx.input_params = list(args[length:]) 148 | with th.no_grad(): 149 | output_tensors = ctx.run_function(*ctx.input_tensors) 150 | return output_tensors 151 | 152 | @staticmethod 153 | def backward(ctx, *output_grads): 154 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 155 | with th.enable_grad(): 156 | # Fixes a bug where the first op in run_function modifies the 157 | # Tensor storage in place, which is not allowed for detach()'d 158 | # Tensors. 159 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 160 | output_tensors = ctx.run_function(*shallow_copies) 161 | input_grads = th.autograd.grad( 162 | output_tensors, 163 | ctx.input_tensors + ctx.input_params, 164 | output_grads, 165 | allow_unused=True, 166 | ) 167 | del ctx.input_tensors 168 | del ctx.input_params 169 | del output_tensors 170 | return (None, None) + input_grads 171 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/reco_train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch_radon import RadonFanbeam, Radon 3 | import numpy as np 4 | import argparse 5 | from pathlib import Path 6 | from .helper import load_tiff_stack_with_metadata, save_to_tiff_stack 7 | 8 | 9 | def para_prepare_1(index): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--path_proj', type=str, default='/root/DDDM/fan_projections.tif', help='Local path of fan beam projection data.') 12 | parser.add_argument('--image_size', type=int, default=736, help='Size of reconstructed image.') 13 | parser.add_argument('--voxel_size', type=float, default=0.7, help='In-slice voxel size [mm].') 14 | parser.add_argument('--fbp_filter', type=str, default='hann', nargs='?',choices=['ram-lak', 'shepp-logan', 'cosine', 'hamming', 'hann'], help='Filter used for FBP.') 15 | args = parser.parse_args() 16 | 17 | _, metadata = load_tiff_stack_with_metadata(Path(args.path_proj)) 18 | 19 | vox_scaling = 1 / args.voxel_size 20 | angles = np.array(metadata['angles'])[:metadata['rotview']] + (np.pi / 2) 21 | if index == 4: 22 | angles = angles[np.arange(3, 1155, index)] 23 | elif index == 1: 24 | angles = angles[np.arange(0, 1152, index)] 25 | elif index == 16: 26 | angles = angles[np.arange(15, 1167, index)] 27 | elif index == 2.5: 28 | angles = (angles[np.arange(1, 1153, 2)])[np.arange(1,289,1)] 29 | elif index == 2: 30 | angles = angles[np.arange(1, 1153, 2)] 31 | radon = RadonFanbeam(args.image_size, 32 | angles, 33 | source_distance=vox_scaling * metadata['dso'], 34 | det_distance=vox_scaling * metadata['ddo'], 35 | det_count=736, 36 | det_spacing=vox_scaling * metadata['du'], 37 | clip_to_circle=True) 38 | return radon 39 | 40 | 41 | def para_prepare_parallel(index): 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument('--path_proj', type=str, default='/root/autodl-tmp/TASK2/fan_projections.tif', help='Local path of fan beam projection data.') 44 | parser.add_argument('--image_size', type=int, default=736, help='Size of reconstructed image.') 45 | parser.add_argument('--voxel_size', type=float, default=0.7, help='In-slice voxel size [mm].') 46 | parser.add_argument('--fbp_filter', type=str, default='hann', nargs='?',choices=['ram-lak', 'shepp-logan', 'cosine', 'hamming', 'hann'], help='Filter used for FBP.') 47 | args = parser.parse_args() 48 | 49 | _, metadata = load_tiff_stack_with_metadata(Path(args.path_proj)) 50 | 51 | vox_scaling = 1 / args.voxel_size 52 | bias = 0# 103 180 53 | angles = np.array(metadata['angles'])[:metadata['rotview']+bias] + (np.pi / 2) 54 | if index == 1.01: 55 | angles = angles[np.arange(0+bias, 1152+bias, 1)] 56 | elif index == 16.5: 57 | angles = (angles[np.arange(0+bias, 1152+bias, 16)])[np.arange(0,36,1)] 58 | elif index == 8.5: 59 | angles = (angles[np.arange(0+bias, 1152+bias, 16)])[np.arange(0,18,1)] 60 | 61 | # radon = RadonFanbeam(args.image_size, 62 | # angles, 63 | # source_distance=vox_scaling * metadata['dso'], 64 | # det_distance=vox_scaling * metadata['ddo'], 65 | # det_count=736, 66 | # det_spacing=vox_scaling * metadata['du'], 67 | # clip_to_circle=False) 68 | radon = Radon(736, angles=angles, clip_to_circle=True) # det_spacing=vox_scaling * metadata['du'], det_count=736 69 | return radon 70 | 71 | def run_reco(projections, radon): 72 | # projections = projections[:,range_clip,:] 73 | if(len(projections.shape) == 4): 74 | sino = torch.flip(projections, dims=[3]) 75 | elif (len(projections.shape) == 3): 76 | sino = torch.flip(projections, dims=[2]) 77 | elif (len(projections.shape) == 2): 78 | sino = torch.flip(projections, dims=[1]) 79 | 80 | filtered_sinogram = radon.filter_sinogram(sino, filter_name='hann') 81 | fbp = 100 * radon.backprojection(filtered_sinogram) 82 | 83 | return fbp 84 | 85 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | import torch.distributed as dist 6 | 7 | 8 | def create_named_schedule_sampler(name, diffusion): 9 | """ 10 | Create a ScheduleSampler from a library of pre-defined samplers. 11 | 12 | :param name: the name of the sampler. 13 | :param diffusion: the diffusion object to sample for. 14 | """ 15 | if name == "uniform": 16 | return UniformSampler(diffusion) 17 | elif name == "loss-second-moment": 18 | return LossSecondMomentResampler(diffusion) 19 | else: 20 | raise NotImplementedError(f"unknown schedule sampler: {name}") 21 | 22 | 23 | class ScheduleSampler(ABC): 24 | """ 25 | A distribution over timesteps in the diffusion process, intended to reduce 26 | variance of the objective. 27 | 28 | By default, samplers perform unbiased importance sampling, in which the 29 | objective's mean is unchanged. 30 | However, subclasses may override sample() to change how the resampled 31 | terms are reweighted, allowing for actual changes in the objective. 32 | """ 33 | 34 | @abstractmethod 35 | def weights(self): 36 | """ 37 | Get a numpy array of weights, one per diffusion step. 38 | 39 | The weights needn't be normalized, but must be positive. 40 | """ 41 | 42 | def sample(self, batch_size, device): 43 | """ 44 | Importance-sample timesteps for a batch. 45 | 46 | :param batch_size: the number of timesteps. 47 | :param device: the torch device to save to. 48 | :return: a tuple (timesteps, weights): 49 | - timesteps: a tensor of timestep indices. 50 | - weights: a tensor of weights to scale the resulting losses. 51 | """ 52 | w = self.weights() 53 | p = w / np.sum(w) 54 | # print("p:", p) 55 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 56 | indices = th.from_numpy(indices_np).long().to(device) 57 | weights_np = 1 / (len(p) * p[indices_np]) 58 | weights = th.from_numpy(weights_np).float().to(device) 59 | return indices, weights 60 | 61 | 62 | class UniformSampler(ScheduleSampler): 63 | def __init__(self, diffusion): 64 | self.diffusion = diffusion 65 | self._weights = np.ones([diffusion.num_timesteps]) 66 | 67 | def weights(self): 68 | return self._weights 69 | 70 | 71 | class LossAwareSampler(ScheduleSampler): 72 | def update_with_local_losses(self, local_ts, local_losses): 73 | """ 74 | Update the reweighting using losses from a model. 75 | 76 | Call this method from each rank with a batch of timesteps and the 77 | corresponding losses for each of those timesteps. 78 | This method will perform synchronization to make sure all of the ranks 79 | maintain the exact same reweighting. 80 | 81 | :param local_ts: an integer Tensor of timesteps. 82 | :param local_losses: a 1D Tensor of losses. 83 | """ 84 | batch_sizes = [ 85 | th.tensor([0], dtype=th.int32, device=local_ts.device) 86 | for _ in range(dist.get_world_size()) 87 | ] 88 | dist.all_gather( 89 | batch_sizes, 90 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 91 | ) 92 | 93 | # Pad all_gather batches to be the maximum batch size. 94 | batch_sizes = [x.item() for x in batch_sizes] 95 | max_bs = max(batch_sizes) 96 | 97 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 98 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 99 | dist.all_gather(timestep_batches, local_ts) 100 | dist.all_gather(loss_batches, local_losses) 101 | timesteps = [ 102 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 103 | ] 104 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 105 | self.update_with_all_losses(timesteps, losses) 106 | 107 | @abstractmethod 108 | def update_with_all_losses(self, ts, losses): 109 | """ 110 | Update the reweighting using losses from a model. 111 | 112 | Sub-classes should override this method to update the reweighting 113 | using losses from the model. 114 | 115 | This method directly updates the reweighting without synchronizing 116 | between workers. It is called by update_with_local_losses from all 117 | ranks with identical arguments. Thus, it should have deterministic 118 | behavior to maintain state across workers. 119 | 120 | :param ts: a list of int timesteps. 121 | :param losses: a list of float losses, one per timestep. 122 | """ 123 | 124 | 125 | class LossSecondMomentResampler(LossAwareSampler): 126 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 127 | self.diffusion = diffusion 128 | self.history_per_term = history_per_term 129 | self.uniform_prob = uniform_prob 130 | self._loss_history = np.zeros( 131 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 132 | ) 133 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 134 | 135 | def weights(self): 136 | if not self._warmed_up(): 137 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 138 | weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) 139 | weights /= np.sum(weights) 140 | weights *= 1 - self.uniform_prob 141 | weights += self.uniform_prob / len(weights) 142 | return weights 143 | 144 | def update_with_all_losses(self, ts, losses): 145 | for t, loss in zip(ts, losses): 146 | if self._loss_counts[t] == self.history_per_term: 147 | # Shift out the oldest loss term. 148 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 149 | self._loss_history[t, -1] = loss 150 | else: 151 | self._loss_history[t, self._loss_counts[t]] = loss 152 | self._loss_counts[t] += 1 153 | 154 | def _warmed_up(self): 155 | return (self._loss_counts == self.history_per_term).all() 156 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/respace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | 4 | from .gaussian_diffusion import GaussianDiffusion 5 | 6 | 7 | def space_timesteps(num_timesteps, section_counts): 8 | """ 9 | Create a list of timesteps to use from an original diffusion process, 10 | given the number of timesteps we want to take from equally-sized portions 11 | of the original process. 12 | 13 | For example, if there's 300 timesteps and the section counts are [10,15,20] 14 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 15 | are strided to be 15 timesteps, and the final 100 are strided to be 20. 16 | 17 | If the stride is a string starting with "ddim", then the fixed striding 18 | from the DDIM paper is used, and only one section is allowed. 19 | 20 | :param num_timesteps: the number of diffusion steps in the original 21 | process to divide up. 22 | :param section_counts: either a list of numbers, or a string containing 23 | comma-separated numbers, indicating the step count 24 | per section. As a special case, use "ddimN" where N 25 | is a number of steps to use the striding from the 26 | DDIM paper. 27 | :return: a set of diffusion steps from the original process to use. 28 | """ 29 | if isinstance(section_counts, str): 30 | if section_counts.startswith("ddim"): 31 | desired_count = int(section_counts[len("ddim") :]) 32 | for i in range(1, num_timesteps): 33 | if len(range(0, num_timesteps, i)) == desired_count: 34 | return set(range(0, num_timesteps, i)) 35 | raise ValueError( 36 | f"cannot create exactly {num_timesteps} steps with an integer stride" 37 | ) 38 | section_counts = [int(x) for x in section_counts.split(",")] 39 | size_per = num_timesteps // len(section_counts) 40 | extra = num_timesteps % len(section_counts) 41 | start_idx = 0 42 | all_steps = [] 43 | for i, section_count in enumerate(section_counts): 44 | size = size_per + (1 if i < extra else 0) 45 | if size < section_count: 46 | raise ValueError( 47 | f"cannot divide section of {size} steps into {section_count}" 48 | ) 49 | if section_count <= 1: 50 | frac_stride = 1 51 | else: 52 | frac_stride = (size - 1) / (section_count - 1) 53 | cur_idx = 0.0 54 | taken_steps = [] 55 | for _ in range(section_count): 56 | taken_steps.append(start_idx + round(cur_idx)) 57 | cur_idx += frac_stride 58 | all_steps += taken_steps 59 | start_idx += size 60 | return set(all_steps) 61 | 62 | 63 | class SpacedDiffusion(GaussianDiffusion): 64 | """ 65 | A diffusion process which can skip steps in a base diffusion process. 66 | 67 | :param use_timesteps: a collection (sequence or set) of timesteps from the 68 | original diffusion process to retain. 69 | :param kwargs: the kwargs to create the base diffusion process. 70 | """ 71 | 72 | def __init__(self, use_timesteps, **kwargs): 73 | self.use_timesteps = set(use_timesteps) 74 | self.timestep_map = [] 75 | self.original_num_steps = len(kwargs["betas"]) 76 | 77 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa 78 | last_alpha_cumprod = 1.0 79 | new_betas = [] 80 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): 81 | if i in self.use_timesteps: 82 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) 83 | last_alpha_cumprod = alpha_cumprod 84 | self.timestep_map.append(i) 85 | kwargs["betas"] = np.array(new_betas) 86 | super().__init__(**kwargs) 87 | 88 | def p_mean_variance( 89 | self, model, *args, **kwargs 90 | ): # pylint: disable=signature-differs 91 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) 92 | 93 | def training_losses( 94 | self, model, *args, **kwargs 95 | ): # pylint: disable=signature-differs 96 | return super().training_losses(self._wrap_model(model), *args, **kwargs) 97 | 98 | def _wrap_model(self, model): 99 | if isinstance(model, _WrappedModel): 100 | return model 101 | return _WrappedModel( 102 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps 103 | ) 104 | 105 | def _scale_timesteps(self, t): 106 | # Scaling is done by the wrapped model. 107 | return t 108 | 109 | 110 | class _WrappedModel: 111 | def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): 112 | self.model = model 113 | self.timestep_map = timestep_map 114 | self.rescale_timesteps = rescale_timesteps 115 | self.original_num_steps = original_num_steps 116 | 117 | def __call__(self, x, ts, **kwargs): 118 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 119 | new_ts = map_tensor[ts] 120 | if self.rescale_timesteps: 121 | new_ts = new_ts.float() * (1000.0 / self.original_num_steps) 122 | return self.model(x, new_ts, **kwargs) 123 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import inspect 3 | 4 | from . import gaussian_diffusion as gd 5 | from .respace import SpacedDiffusion, space_timesteps 6 | from .unet import SuperResModel, UNetModel 7 | 8 | NUM_CLASSES = 2 9 | 10 | 11 | def model_and_diffusion_defaults(): 12 | """ 13 | Defaults for image training. 14 | """ 15 | return dict( 16 | image_size=512, 17 | num_channels=128, 18 | num_res_blocks=2, 19 | num_heads=4, 20 | num_heads_upsample=-1, 21 | attention_resolutions="32,16", 22 | dropout=0.3, 23 | learn_sigma=False, 24 | sigma_small=False, 25 | class_cond=False, 26 | diffusion_steps=4000, 27 | noise_schedule="cosine", 28 | timestep_respacing="2", 29 | use_kl=False, 30 | predict_xstart=True, 31 | rescale_timesteps=False, 32 | rescale_learned_sigmas=False, 33 | use_checkpoint=False, 34 | use_scale_shift_norm=True, 35 | ) 36 | 37 | 38 | def create_model_and_diffusion( 39 | image_size, 40 | class_cond, 41 | learn_sigma, 42 | sigma_small, 43 | num_channels, 44 | num_res_blocks, 45 | num_heads, 46 | num_heads_upsample, 47 | attention_resolutions, 48 | dropout, 49 | diffusion_steps, 50 | noise_schedule, 51 | timestep_respacing, 52 | use_kl, 53 | predict_xstart, 54 | rescale_timesteps, 55 | rescale_learned_sigmas, 56 | use_checkpoint, 57 | use_scale_shift_norm, 58 | ): 59 | model = create_model( 60 | image_size, 61 | num_channels, 62 | num_res_blocks, 63 | learn_sigma=learn_sigma, 64 | class_cond=class_cond, 65 | use_checkpoint=use_checkpoint, 66 | attention_resolutions=attention_resolutions, 67 | num_heads=num_heads, 68 | num_heads_upsample=num_heads_upsample, 69 | use_scale_shift_norm=use_scale_shift_norm, 70 | dropout=dropout, 71 | ) 72 | diffusion = create_gaussian_diffusion( 73 | steps=diffusion_steps, 74 | learn_sigma=learn_sigma, 75 | sigma_small=sigma_small, 76 | noise_schedule=noise_schedule, 77 | use_kl=use_kl, 78 | predict_xstart=predict_xstart, 79 | rescale_timesteps=rescale_timesteps, 80 | rescale_learned_sigmas=rescale_learned_sigmas, 81 | timestep_respacing=timestep_respacing, 82 | ) 83 | return model, diffusion 84 | 85 | 86 | def create_model( 87 | image_size, 88 | num_channels, 89 | num_res_blocks, 90 | learn_sigma, 91 | class_cond, 92 | use_checkpoint, 93 | attention_resolutions, 94 | num_heads, 95 | num_heads_upsample, 96 | use_scale_shift_norm, 97 | dropout, 98 | ): 99 | if image_size == 256 or 512: 100 | channel_mult = (1, 1, 2, 2, 4, 4) 101 | print("channel_mult like this") 102 | elif image_size == 64: 103 | channel_mult = (1, 2, 3, 4) 104 | elif image_size == 32: 105 | channel_mult = (1, 2, 2, 2) 106 | else: 107 | raise ValueError(f"unsupported image size: {image_size}") 108 | 109 | attention_ds = [] 110 | for res in attention_resolutions.split(","): 111 | attention_ds.append(image_size // int(res)) 112 | 113 | return UNetModel( 114 | in_channels=1, 115 | model_channels=num_channels, 116 | out_channels=(1 if not learn_sigma else 2), 117 | num_res_blocks=num_res_blocks, 118 | attention_resolutions=tuple(attention_ds), 119 | dropout=dropout, 120 | channel_mult=channel_mult, 121 | num_classes=(NUM_CLASSES if class_cond else None), 122 | use_checkpoint=use_checkpoint, 123 | num_heads=num_heads, 124 | num_heads_upsample=num_heads_upsample, 125 | use_scale_shift_norm=use_scale_shift_norm, 126 | ) 127 | 128 | 129 | def sr_model_and_diffusion_defaults(): 130 | res = model_and_diffusion_defaults() 131 | print("timestep_respacing:", res["timestep_respacing"]) 132 | res["large_size"] = 512 133 | res["small_size"] = [512, 512] 134 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 135 | for k in res.copy().keys(): 136 | if k not in arg_names: 137 | del res[k] 138 | return res 139 | 140 | 141 | def sr_create_model_and_diffusion( 142 | large_size, 143 | small_size, 144 | class_cond, 145 | learn_sigma, 146 | num_channels, 147 | num_res_blocks, 148 | num_heads, 149 | num_heads_upsample, 150 | attention_resolutions, 151 | dropout, 152 | diffusion_steps, 153 | noise_schedule, 154 | timestep_respacing, 155 | use_kl, 156 | predict_xstart, 157 | rescale_timesteps, 158 | rescale_learned_sigmas, 159 | use_checkpoint, 160 | use_scale_shift_norm, 161 | ): 162 | model = sr_create_model( 163 | large_size, 164 | small_size, 165 | num_channels, 166 | num_res_blocks, 167 | learn_sigma=learn_sigma, 168 | class_cond=class_cond, 169 | use_checkpoint=use_checkpoint, 170 | attention_resolutions=attention_resolutions, 171 | num_heads=num_heads, 172 | num_heads_upsample=num_heads_upsample, 173 | use_scale_shift_norm=use_scale_shift_norm, 174 | dropout=dropout, 175 | ) 176 | diffusion = create_gaussian_diffusion( 177 | steps=diffusion_steps, 178 | learn_sigma=learn_sigma, 179 | noise_schedule=noise_schedule, 180 | use_kl=use_kl, 181 | predict_xstart=predict_xstart, 182 | rescale_timesteps=rescale_timesteps, 183 | rescale_learned_sigmas=rescale_learned_sigmas, 184 | timestep_respacing=timestep_respacing, 185 | ) 186 | return model, diffusion 187 | 188 | 189 | def sr_create_model( 190 | large_size, 191 | small_size, 192 | num_channels, 193 | num_res_blocks, 194 | learn_sigma, 195 | class_cond, 196 | use_checkpoint, 197 | attention_resolutions, 198 | num_heads, 199 | num_heads_upsample, 200 | use_scale_shift_norm, 201 | dropout, 202 | ): 203 | _ = small_size # hack to prevent unused variable 204 | 205 | if large_size == 256 or 512: 206 | channel_mult = (1, 1, 2, 2, 4, 4) 207 | print("channel_mult like this 111111") 208 | elif large_size == 64: 209 | channel_mult = (1, 2, 3, 4) 210 | else: 211 | raise ValueError(f"unsupported large size: {large_size}") 212 | 213 | attention_ds = [] 214 | for res in attention_resolutions.split(","): 215 | attention_ds.append(large_size // int(res)) 216 | 217 | return SuperResModel( 218 | in_channels=1, 219 | model_channels=num_channels, 220 | out_channels=(1 if not learn_sigma else 2), 221 | num_res_blocks=num_res_blocks, 222 | attention_resolutions=tuple(attention_ds), 223 | dropout=dropout, 224 | channel_mult=channel_mult, 225 | num_classes=(NUM_CLASSES if class_cond else None), 226 | use_checkpoint=use_checkpoint, 227 | num_heads=num_heads, 228 | num_heads_upsample=num_heads_upsample, 229 | use_scale_shift_norm=use_scale_shift_norm, 230 | ) 231 | 232 | 233 | def create_gaussian_diffusion( 234 | *, 235 | steps=4000, 236 | learn_sigma=False, 237 | sigma_small=False, 238 | noise_schedule="cosine", 239 | use_kl=False, 240 | predict_xstart=True, 241 | rescale_timesteps=False, 242 | rescale_learned_sigmas=False, 243 | timestep_respacing="", 244 | ): 245 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 246 | if use_kl: 247 | loss_type = gd.LossType.RESCALED_KL 248 | elif rescale_learned_sigmas: 249 | loss_type = gd.LossType.RESCALED_MSE 250 | else: 251 | loss_type = gd.LossType.MSE 252 | if not timestep_respacing: 253 | timestep_respacing = [steps] 254 | return SpacedDiffusion( 255 | use_timesteps=space_timesteps(steps, timestep_respacing), 256 | betas=betas, 257 | model_mean_type=( 258 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 259 | ), 260 | model_var_type=( 261 | ( 262 | gd.ModelVarType.FIXED_LARGE 263 | if not sigma_small 264 | else gd.ModelVarType.FIXED_SMALL 265 | ) 266 | if not learn_sigma 267 | else gd.ModelVarType.LEARNED_RANGE 268 | ), 269 | loss_type=loss_type, 270 | rescale_timesteps=rescale_timesteps, 271 | ) 272 | 273 | 274 | def add_dict_to_argparser(parser, default_dict): 275 | for k, v in default_dict.items(): 276 | v_type = type(v) 277 | if v is None: 278 | v_type = str 279 | elif isinstance(v, bool): 280 | v_type = str2bool 281 | parser.add_argument(f"--{k}", default=v, type=v_type) 282 | 283 | 284 | def args_to_dict(args, keys): 285 | return {k: getattr(args, k) for k in keys} 286 | 287 | 288 | def str2bool(v): 289 | """ 290 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 291 | """ 292 | if isinstance(v, bool): 293 | return v 294 | if v.lower() in ("yes", "true", "t", "y", "1"): 295 | return True 296 | elif v.lower() in ("no", "false", "f", "n", "0"): 297 | return False 298 | else: 299 | raise argparse.ArgumentTypeError("boolean value expected") 300 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/scripts/image_nll.py: -------------------------------------------------------------------------------- 1 | """ 2 | Approximate the bits/dimension for an image model. 3 | """ 4 | 5 | import argparse 6 | import os 7 | 8 | import numpy as np 9 | import torch.distributed as dist 10 | 11 | from improved_diffusion import dist_util, logger 12 | from improved_diffusion.image_datasets import load_data 13 | from improved_diffusion.script_util import ( 14 | model_and_diffusion_defaults, 15 | create_model_and_diffusion, 16 | add_dict_to_argparser, 17 | args_to_dict, 18 | ) 19 | 20 | 21 | def main(): 22 | args = create_argparser().parse_args() 23 | 24 | dist_util.setup_dist() 25 | logger.configure() 26 | 27 | logger.log("creating model and diffusion...") 28 | model, diffusion = create_model_and_diffusion( 29 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 30 | ) 31 | model.load_state_dict( 32 | dist_util.load_state_dict(args.model_path, map_location="cpu") 33 | ) 34 | model.to(dist_util.dev()) 35 | model.eval() 36 | 37 | logger.log("creating data loader...") 38 | data = load_data( 39 | data_dir=args.data_dir, 40 | batch_size=args.batch_size, 41 | image_size=args.image_size, 42 | class_cond=args.class_cond, 43 | deterministic=True, 44 | ) 45 | 46 | logger.log("evaluating...") 47 | run_bpd_evaluation(model, diffusion, data, args.num_samples, args.clip_denoised) 48 | 49 | 50 | def run_bpd_evaluation(model, diffusion, data, num_samples, clip_denoised): 51 | all_bpd = [] 52 | all_metrics = {"vb": [], "mse": [], "xstart_mse": []} 53 | num_complete = 0 54 | while num_complete < num_samples: 55 | batch, model_kwargs = next(data) 56 | batch = batch.to(dist_util.dev()) 57 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 58 | minibatch_metrics = diffusion.calc_bpd_loop( 59 | model, batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs 60 | ) 61 | 62 | for key, term_list in all_metrics.items(): 63 | terms = minibatch_metrics[key].mean(dim=0) / dist.get_world_size() 64 | dist.all_reduce(terms) 65 | term_list.append(terms.detach().cpu().numpy()) 66 | 67 | total_bpd = minibatch_metrics["total_bpd"] 68 | total_bpd = total_bpd.mean() / dist.get_world_size() 69 | dist.all_reduce(total_bpd) 70 | all_bpd.append(total_bpd.item()) 71 | num_complete += dist.get_world_size() * batch.shape[0] 72 | 73 | logger.log(f"done {num_complete} samples: bpd={np.mean(all_bpd)}") 74 | 75 | if dist.get_rank() == 0: 76 | for name, terms in all_metrics.items(): 77 | out_path = os.path.join(logger.get_dir(), f"{name}_terms.npz") 78 | logger.log(f"saving {name} terms to {out_path}") 79 | np.savez(out_path, np.mean(np.stack(terms), axis=0)) 80 | 81 | dist.barrier() 82 | logger.log("evaluation complete") 83 | 84 | 85 | def create_argparser(): 86 | defaults = dict( 87 | data_dir="", clip_denoised=True, num_samples=1000, batch_size=1, model_path="" 88 | ) 89 | defaults.update(model_and_diffusion_defaults()) 90 | parser = argparse.ArgumentParser() 91 | add_dict_to_argparser(parser, defaults) 92 | return parser 93 | 94 | 95 | if __name__ == "__main__": 96 | main() 97 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/scripts/img_multi_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of samples from a super resolution model, given a batch 3 | of samples from a regular model from image_sample.py. 4 | """ 5 | import matplotlib.pyplot as plt 6 | import argparse 7 | import os 8 | import torch.nn.functional as F 9 | from improved_diffusion.image_datasets import trans_data_preprocess 10 | import blobfile as bf 11 | import numpy as np 12 | import torch as th 13 | import torch.distributed as dist 14 | from skimage.metrics import structural_similarity as ssim 15 | from skimage.metrics import peak_signal_noise_ratio as psnr 16 | import lpips 17 | from improved_diffusion.reco_train import run_reco, para_prepare_parallel 18 | import PIL.Image as Image 19 | 20 | 21 | from improved_diffusion import dist_util, logger 22 | from improved_diffusion.script_util import ( 23 | sr_model_and_diffusion_defaults, 24 | sr_create_model_and_diffusion, 25 | args_to_dict, 26 | add_dict_to_argparser, 27 | ) 28 | 29 | model_name = "ema_0.9999_062500" 30 | 31 | def main(): 32 | num = 800 # total image numbers 33 | 34 | args = create_argparser().parse_args() 35 | 36 | dist_util.setup_dist() 37 | logger.configure() 38 | 39 | logger.log("creating model...") 40 | model, diffusion = sr_create_model_and_diffusion( 41 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 42 | ) 43 | model.load_state_dict( 44 | dist_util.load_state_dict(args.model_path, map_location="cpu") 45 | ) 46 | model.to(dist_util.dev()) 47 | model.eval() 48 | 49 | logger.log("loading data...") 50 | 51 | radon_36 = para_prepare_parallel(16.5) 52 | radon_72 = para_prepare_parallel(8.5) 53 | radon_1152 = para_prepare_parallel(1.01) 54 | data = load_superres_data(args.batch_size, radon_36=radon_36, radon_1152=radon_1152) 55 | 56 | 57 | logger.log("creating samples...") 58 | 59 | MSE, SSIM, PSNR, LPIP = [], [], [], [] 60 | RAW, RECON = [], [] 61 | MSE_36, SSIM_36, PSNR_36, LPIP_36 = [], [], [], [] 62 | lpip_loss = lpips.LPIPS(net="alex").to(dist_util.dev()) 63 | l2loss = th.nn.MSELoss().to(dist_util.dev()) 64 | 65 | 66 | 67 | for i in range(0, num//args.batch_size):# 68 | raw_img, model_kwargs = next(data) 69 | index = model_kwargs.pop('index') 70 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 71 | 72 | 73 | sample_fn = ( 74 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 75 | ) 76 | sample = sample_fn( 77 | model, 78 | (args.batch_size, 1, 512, 512), 79 | clip_denoised=args.clip_denoised, 80 | model_kwargs=model_kwargs, 81 | ) 82 | 83 | sample = sample + 10 # make sure samples are all positive 84 | raw_img = raw_img + 10 85 | 86 | 87 | model_kwargs["_36_res"] = model_kwargs["_36_res"] + 10 88 | 89 | assert sample.min()>0 and raw_img.min()>0 90 | output_fbp_npy = sample.cpu().detach().numpy().squeeze() 91 | RAW.append(th.mean(raw_img).item()) 92 | RECON.append(th.mean(sample).item()) 93 | 94 | for j in range(0, args.batch_size): 95 | 96 | raw_npy = raw_img.cpu().detach().numpy() 97 | 98 | l2loss_value = l2loss(sample[j], raw_img[j]).item() 99 | print("index:", index[j].item(), "MSELoss:", l2loss_value) 100 | MSE.append(l2loss_value) 101 | 102 | ssim_value = ssim(np.squeeze(output_fbp_npy[j]),np.squeeze(raw_npy[j]), data_range = raw_npy[j].max() - raw_npy[j].min()) 103 | print("index:", index[j].item(), "SSIM:", ssim_value) 104 | SSIM.append(ssim_value) 105 | psnr_value = psnr(np.squeeze(output_fbp_npy[j]),np.squeeze(raw_npy[j]), data_range = raw_npy[j].max() - raw_npy[j].min()) 106 | PSNR.append(psnr_value) 107 | 108 | lpip_value = lpip_loss(sample[j] - 10, raw_img[j] - 10) 109 | print("lpips:", lpip_value.item()) 110 | LPIP.append(lpip_value.item()) 111 | 112 | ssim_value_36 = ssim((model_kwargs["_36_res"][j]).cpu().detach().numpy().squeeze(),np.squeeze(raw_npy[j]),data_range = raw_npy[j].max() - raw_npy[j].min()) # [:,100:412,100:412] 113 | SSIM_36.append(ssim_value_36) 114 | l2loss_value_36 = l2loss((model_kwargs["_36_res"][j]), raw_img[j]).item() 115 | MSE_36.append(l2loss_value_36) 116 | psnr_value_36 = psnr((model_kwargs["_36_res"][j]).cpu().detach().numpy().squeeze(),np.squeeze(raw_npy[j]),data_range = raw_npy[j].max() - raw_npy[j].min()) 117 | PSNR_36.append(psnr_value_36) 118 | lpip_value_36 = lpip_loss((model_kwargs["_36_res"][j]) - 10, raw_img[j] - 10) 119 | LPIP_36.append(lpip_value_36.item()) 120 | 121 | 122 | print("mse mean:", np.mean(MSE)) 123 | print("ssim mean:", np.mean(SSIM)) 124 | print("lpip mean:", np.mean(LPIP)) 125 | print("psnr mean:", np.mean(PSNR)) 126 | print("raw mean:", np.mean(RAW)) 127 | print("recon mean:", np.mean(RECON)) 128 | 129 | print("mse mean:", np.mean(MSE_36)) # metrics for original 36-view CT image 130 | print("ssim mean:", np.mean(SSIM_36)) 131 | print("lpip mean:", np.mean(LPIP_36)) 132 | print("psnr mean:", np.mean(PSNR_36)) 133 | 134 | 135 | def create_argparser(): 136 | defaults = dict( 137 | clip_denoised=False, 138 | num_samples=1, 139 | batch_size=10, 140 | use_ddim=True, 141 | base_samples="", 142 | model_path=f'/root/autodl-tmp/IRM/improved_diffusion/model_save_3channels/{model_name}.pt', 143 | ) 144 | print(defaults["model_path"]) 145 | defaults.update(sr_model_and_diffusion_defaults()) 146 | parser = argparse.ArgumentParser() 147 | add_dict_to_argparser(parser, defaults) 148 | return parser 149 | 150 | 151 | def norm(fbp_img): 152 | return (fbp_img - (-0.75)) 153 | 154 | def load_superres_data(batch_size, radon_36, radon_1152): 155 | 156 | data = trans_data_preprocess(batch_size = batch_size, shuffle = False, num_workers = 8) 157 | 158 | for large_batch, model_kwargs in data: 159 | # large_batch = norm(large_batch.to("cuda") - 1.) 160 | large_batch = norm(th.flip(run_reco((large_batch).to("cuda"), radon_1152) - 1.,dims=[2,3])[:,:,112:624,112:624]) 161 | model_kwargs["low_res"] = norm(model_kwargs["low_res"].to("cuda") - 1.) 162 | model_kwargs["_36_res"] = norm(th.flip(run_reco(model_kwargs["_36_res"][:,:,np.arange(0,576,8),:].to("cuda"), radon_36) - 1.,dims=[2,3])[:,:,112:624,112:624]) 163 | _36_npy = model_kwargs["_36_res"].cpu().detach().numpy().squeeze() 164 | yield large_batch, model_kwargs 165 | 166 | 167 | if __name__ == "__main__": 168 | main() 169 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/scripts/img_super_res_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of samples from a super resolution model, given a batch 3 | of samples from a regular model from image_sample.py. 4 | """ 5 | import matplotlib.pyplot as plt 6 | import argparse 7 | import os 8 | import torch.nn.functional as F 9 | from improved_diffusion.image_datasets import improved_data_preprocess_val 10 | import blobfile as bf 11 | import numpy as np 12 | import torch as th 13 | import torch.distributed as dist 14 | from skimage.metrics import structural_similarity as ssim 15 | from skimage.metrics import peak_signal_noise_ratio as psnr 16 | import lpips 17 | from pytorch_msssim import ms_ssim, MS_SSIM, SSIM 18 | 19 | 20 | from improved_diffusion import dist_util, logger 21 | from improved_diffusion.script_util import ( 22 | sr_model_and_diffusion_defaults, 23 | sr_create_model_and_diffusion, 24 | args_to_dict, 25 | add_dict_to_argparser, 26 | ) 27 | from improved_diffusion.train_util import TrainLoop 28 | from improved_diffusion.reco_train import run_reco, para_prepare_parallel 29 | 30 | 31 | index = 1 32 | model_name = "model100000" 33 | 34 | def main(): 35 | args = create_argparser().parse_args() 36 | 37 | dist_util.setup_dist() 38 | logger.configure() 39 | 40 | logger.log("creating model...") 41 | model, diffusion = sr_create_model_and_diffusion( 42 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 43 | ) 44 | model.load_state_dict( 45 | dist_util.load_state_dict(args.model_path, map_location="cpu") 46 | ) 47 | model.to(dist_util.dev()) 48 | model.eval() 49 | 50 | radon_72 = para_prepare_parallel(16.5) 51 | radon_1152 = para_prepare_parallel(1.01) 52 | 53 | 54 | logger.log("loading data...") 55 | data = load_superres_data( 56 | args.batch_size, 57 | large_size=args.large_size, 58 | small_size=args.small_size, 59 | class_cond=args.class_cond, 60 | radon_72=radon_72, 61 | radon_1152=radon_1152 62 | ) 63 | 64 | logger.log("creating samples...") 65 | all_images = [] 66 | while len(all_images) * args.batch_size < args.num_samples: 67 | raw_img, model_kwargs = next(data) 68 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 69 | 70 | npy_input = model_kwargs['low_res'].squeeze().cpu().detach().numpy() 71 | np.save(f"/root/autodl-tmp/TASK7/improved_diffusion/test_pic/{index}_{model_name}_input.npy", npy_input) 72 | plt.imshow(npy_input, cmap=plt.cm.bone) 73 | plt.savefig(f"/root/autodl-tmp/TASK7/improved_diffusion/test_pic/{index}_{model_name}_input.png") 74 | 75 | 76 | sample_fn = ( 77 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 78 | ) 79 | 80 | 81 | sample = sample_fn( 82 | model, 83 | (args.batch_size, 1, 512, 512), #args.large_size, args.large_size 84 | clip_denoised=args.clip_denoised, 85 | model_kwargs=model_kwargs, 86 | ) 87 | 88 | npy = sample.squeeze().cpu().detach().numpy() 89 | 90 | plt.imshow(npy, cmap=plt.cm.bone) 91 | 92 | l2loss = th.nn.MSELoss().to(dist_util.dev()) 93 | print("MSELoss:", l2loss(sample, raw_img.to(dist_util.dev())).item()) 94 | raw_npy = raw_img.squeeze().cpu().detach().numpy() 95 | plt.imshow(raw_npy, cmap=plt.cm.bone) 96 | np.save(f"/root/autodl-tmp/TASK7/improved_diffusion/test_pic/{index}_{model_name}_raw.npy", raw_npy) 97 | plt.savefig(f"/root/autodl-tmp/TASK7/improved_diffusion/test_pic/{index}_{model_name}_raw.png") 98 | print("PSNR:", psnr(npy, raw_npy, data_range=np.max(raw_npy)-np.min(raw_npy))) 99 | 100 | _ms_ssim = SSIM(win_size=7, win_sigma=1.5, data_range=1, size_average=False, channel=1) 101 | print("MS_SSIM:", _ms_ssim((sample+1)/4, ((raw_img+1)/4).to(dist_util.dev()))) 102 | print("SSIM:", ssim(npy, raw_npy, data_range=raw_npy.max()-raw_npy.min())) # 103 | lpip_loss = lpips.LPIPS(net="alex").to(dist_util.dev()) 104 | 105 | print("max sample:", th.max(sample).item()) 106 | print("min sample:", th.min(sample).item()) 107 | print("max raw_img:", th.max(raw_img).item()) 108 | print("min raw_img:", th.min(raw_img).item()) 109 | lpip_value = lpip_loss(sample, raw_img.to(dist_util.dev())) 110 | print("lpips:", lpip_value.item()) 111 | 112 | 113 | print("MSELoss input:", l2loss(model_kwargs['low_res'], raw_img.to(dist_util.dev())).item()) 114 | print("SSIM input:", ssim(npy_input, raw_npy, data_range=raw_npy.max()-raw_npy.min())) 115 | print("PSNR input:", psnr(npy_input, raw_npy, data_range=raw_npy.max()-raw_npy.min())) 116 | lpip_value = lpip_loss(model_kwargs['low_res'], raw_img.to(dist_util.dev())) 117 | print("lpips input:", lpip_value.item()) 118 | 119 | break 120 | 121 | 122 | 123 | 124 | def create_argparser(): 125 | defaults = dict( 126 | clip_denoised=True, 127 | num_samples=1, 128 | batch_size=1, 129 | use_ddim=True, 130 | base_samples="", 131 | model_path=f'/root/autodl-tmp/TASK7/improved_diffusion/model_save_noise/{model_name}.pt', 132 | ) 133 | print(defaults["model_path"]) 134 | defaults.update(sr_model_and_diffusion_defaults()) 135 | parser = argparse.ArgumentParser() 136 | add_dict_to_argparser(parser, defaults) 137 | return parser 138 | 139 | def norm(fbp_img): 140 | return ((fbp_img - (-0.75)) / 1.5) #-0.75 is the average value of whole CT images, 1.5 is the factor that make black area linearly mapping to -1 141 | 142 | 143 | def load_superres_data(batch_size, large_size, small_size, class_cond, radon_72, radon_1152): 144 | print("class_cond:", class_cond) 145 | data = improved_data_preprocess_val(batch_size = batch_size, shuffle = True, num_workers = 8) 146 | for large_batch, model_kwargs in data: 147 | # large_batch = norm(large_batch.to("cuda") - 1.) 148 | large_batch = norm(large_batch.to("cuda") - 1.) 149 | model_kwargs["low_res"] = norm(model_kwargs["low_res"].to("cuda") - 1.) 150 | # model_kwargs["_72_res"] = norm(torch.flip(run_reco(F.interpolate((model_kwargs["_72_res"][:,:,np.arange(0,576,2),:] ).to("cuda"), (72, 736), mode="nearest"), radon_72) - 1.,dims=[2,3])[:,:,112:624,112:624]) 151 | 152 | yield large_batch, model_kwargs 153 | 154 | if __name__ == "__main__": 155 | main() 156 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/scripts/img_super_res_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a super-resolution model. 3 | """ 4 | 5 | import argparse 6 | 7 | import torch.nn.functional as F 8 | import torch 9 | import numpy as np 10 | 11 | from improved_diffusion import dist_util, logger 12 | from improved_diffusion.image_datasets import improved_data_preprocess 13 | from improved_diffusion.resample import create_named_schedule_sampler 14 | from improved_diffusion.script_util import ( 15 | sr_model_and_diffusion_defaults, 16 | sr_create_model_and_diffusion, 17 | args_to_dict, 18 | add_dict_to_argparser, 19 | ) 20 | from improved_diffusion.train_util import TrainLoop 21 | from improved_diffusion.reco_train import run_reco, para_prepare_parallel 22 | 23 | 24 | def main(): 25 | args = create_argparser().parse_args() 26 | 27 | dist_util.setup_dist() 28 | logger.configure() 29 | 30 | logger.log("creating model...") 31 | model, diffusion = sr_create_model_and_diffusion( 32 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 33 | ) 34 | model.to(dist_util.dev()) 35 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 36 | 37 | logger.log("creating data loader...") 38 | 39 | radon_36 = para_prepare_parallel(16.5) 40 | radon_1152 = para_prepare_parallel(1.01) 41 | 42 | data = load_superres_data( 43 | args.batch_size, 44 | large_size=args.large_size, 45 | small_size=args.small_size, 46 | class_cond=args.class_cond, 47 | radon_36=radon_36, 48 | radon_1152=radon_1152 49 | ) 50 | 51 | batch, cond = next(data) 52 | print(batch.shape) 53 | print(cond["low_res"].shape) 54 | 55 | for i in range(0,10): 56 | batch, cond = next(data) 57 | print(torch.max(batch[0])) 58 | print(torch.min(batch[0])) 59 | print(torch.mean(batch[0])) 60 | print(torch.max(cond["low_res"][0])) 61 | print(torch.min(cond["low_res"][0])) 62 | print(torch.mean(cond["low_res"][0])) 63 | print(" ") 64 | 65 | 66 | logger.log("training...") 67 | print(args) 68 | TrainLoop( 69 | model=model, 70 | diffusion=diffusion, 71 | data=data, 72 | batch_size=args.batch_size, 73 | microbatch=args.microbatch, 74 | lr=args.lr, 75 | ema_rate=args.ema_rate, 76 | log_interval=args.log_interval, 77 | save_interval=args.save_interval, 78 | resume_checkpoint=args.resume_checkpoint, 79 | use_fp16=args.use_fp16, 80 | fp16_scale_growth=args.fp16_scale_growth, 81 | schedule_sampler=schedule_sampler, 82 | weight_decay=args.weight_decay, 83 | lr_anneal_steps=args.lr_anneal_steps, 84 | ).run_loop() 85 | 86 | def norm(fbp_img): 87 | return (fbp_img - (-0.75)) 88 | 89 | def load_superres_data(batch_size, large_size, small_size, class_cond, radon_36, radon_1152): 90 | print("class_cond:", class_cond) 91 | data = improved_data_preprocess(batch_size = batch_size, shuffle = True, num_workers = 8) 92 | for large_batch, model_kwargs in data: 93 | large_batch = norm(large_batch.to("cuda") - 1.) 94 | model_kwargs["low_res"] = norm(model_kwargs["low_res"].to("cuda") - 1.) 95 | model_kwargs["_36_res"] = norm(torch.flip(run_reco(model_kwargs["_36_res"][:,:,np.arange(0,576,16),:].to("cuda"), radon_36) - 1.,dims=[2,3])[:,:,112:624,112:624]) 96 | 97 | yield large_batch, model_kwargs 98 | 99 | 100 | def create_argparser(): 101 | defaults = dict( 102 | data_dir="", 103 | schedule_sampler="loss-second-moment", 104 | lr=1e-4, 105 | weight_decay=0.0, 106 | lr_anneal_steps=0, 107 | batch_size=2, 108 | microbatch=-1, 109 | ema_rate="0.9999", 110 | log_interval=250, 111 | save_interval=2500, 112 | resume_checkpoint="/root/autodl-tmp/TASK7/improved_diffusion/model_save_3channels_modi/model010000.pt",# 113 | use_fp16=False, 114 | fp16_scale_growth=1e-3, 115 | ) 116 | defaults.update(sr_model_and_diffusion_defaults()) 117 | parser = argparse.ArgumentParser() 118 | add_dict_to_argparser(parser, defaults) 119 | return parser 120 | 121 | 122 | if __name__ == "__main__": 123 | main() 124 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import numpy as np 7 | import torch as th 8 | import torch.distributed as dist 9 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 10 | from torch.optim import AdamW 11 | import lpips 12 | from pytorch_msssim import ms_ssim, MS_SSIM, SSIM 13 | 14 | from . import dist_util, logger 15 | from .fp16_util import ( 16 | make_master_params, 17 | master_params_to_model_params, 18 | model_grads_to_master_grads, 19 | unflatten_master_params, 20 | zero_grad, 21 | ) 22 | from .nn import update_ema 23 | from .resample import LossAwareSampler, UniformSampler 24 | 25 | # For ImageNet experiments, this was a good default value. 26 | # We found that the lg_loss_scale quickly climbed to 27 | # 20-21 within the first ~1K steps of training. 28 | INITIAL_LOG_LOSS_SCALE = 20.0 29 | 30 | 31 | class TrainLoop: 32 | def __init__( 33 | self, 34 | *, 35 | model, 36 | diffusion, 37 | data, 38 | batch_size, 39 | microbatch, 40 | lr, 41 | ema_rate, 42 | log_interval, 43 | save_interval, 44 | resume_checkpoint, 45 | use_fp16=False, 46 | fp16_scale_growth=1e-3, 47 | schedule_sampler=None, 48 | weight_decay=0.0, 49 | lr_anneal_steps=0, 50 | ): 51 | self.model = model 52 | self.diffusion = diffusion 53 | self.data = data 54 | self.batch_size = batch_size 55 | self.microbatch = microbatch if microbatch > 0 else batch_size 56 | self.lr = lr 57 | self.ema_rate = ( 58 | [ema_rate] 59 | if isinstance(ema_rate, float) 60 | else [float(x) for x in ema_rate.split(",")] 61 | ) 62 | self.log_interval = log_interval 63 | self.save_interval = save_interval 64 | self.resume_checkpoint = resume_checkpoint 65 | self.use_fp16 = use_fp16 66 | self.fp16_scale_growth = fp16_scale_growth 67 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 68 | self.weight_decay = weight_decay 69 | self.lr_anneal_steps = lr_anneal_steps 70 | 71 | self.step = 0 72 | self.resume_step = 0 73 | self.global_batch = self.batch_size * dist.get_world_size() 74 | 75 | self.model_params = list(self.model.parameters()) 76 | self.master_params = self.model_params 77 | self.lg_loss_scale = INITIAL_LOG_LOSS_SCALE 78 | self.sync_cuda = th.cuda.is_available() 79 | 80 | self._load_and_sync_parameters() 81 | if self.use_fp16: 82 | self._setup_fp16() 83 | 84 | self.opt = AdamW(self.master_params, lr=self.lr, weight_decay=self.weight_decay) 85 | if self.resume_step: 86 | self._load_optimizer_state() 87 | # Model was resumed, either due to a restart or a checkpoint 88 | # being specified at the command line. 89 | self.ema_params = [ 90 | self._load_ema_parameters(rate) for rate in self.ema_rate 91 | ] 92 | else: 93 | self.ema_params = [ 94 | copy.deepcopy(self.master_params) for _ in range(len(self.ema_rate)) 95 | ] 96 | 97 | if th.cuda.is_available(): 98 | self.use_ddp = True 99 | self.ddp_model = DDP( 100 | self.model, 101 | device_ids=[dist_util.dev()], 102 | output_device=dist_util.dev(), 103 | broadcast_buffers=False, 104 | bucket_cap_mb=128, 105 | find_unused_parameters=False, 106 | ) 107 | else: 108 | if dist.get_world_size() > 1: 109 | logger.warn( 110 | "Distributed training requires CUDA. " 111 | "Gradients will not be synchronized properly!" 112 | ) 113 | self.use_ddp = False 114 | self.ddp_model = self.model 115 | 116 | def _load_and_sync_parameters(self): 117 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 118 | 119 | if resume_checkpoint: 120 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 121 | if dist.get_rank() == 0: 122 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 123 | self.model.load_state_dict( 124 | dist_util.load_state_dict( 125 | resume_checkpoint, map_location=dist_util.dev() 126 | ) 127 | ) 128 | 129 | dist_util.sync_params(self.model.parameters()) 130 | 131 | def _load_ema_parameters(self, rate): 132 | ema_params = copy.deepcopy(self.master_params) 133 | 134 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 135 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 136 | if ema_checkpoint: 137 | if dist.get_rank() == 0: 138 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 139 | state_dict = dist_util.load_state_dict( 140 | ema_checkpoint, map_location=dist_util.dev() 141 | ) 142 | ema_params = self._state_dict_to_master_params(state_dict) 143 | 144 | dist_util.sync_params(ema_params) 145 | return ema_params 146 | 147 | def _load_optimizer_state(self): 148 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 149 | opt_checkpoint = bf.join( 150 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 151 | ) 152 | if bf.exists(opt_checkpoint): 153 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 154 | state_dict = dist_util.load_state_dict( 155 | opt_checkpoint, map_location=dist_util.dev() 156 | ) 157 | self.opt.load_state_dict(state_dict) 158 | 159 | def _setup_fp16(self): 160 | self.master_params = make_master_params(self.model_params) 161 | self.model.convert_to_fp16() 162 | 163 | def run_loop(self): 164 | lpip_fn = lpips.LPIPS(net="alex").to(dist_util.dev()) ######## 165 | ssim_fn = MS_SSIM(win_size=7, win_sigma=1.5, data_range=1, size_average=False, channel=1) 166 | while ( 167 | not self.lr_anneal_steps 168 | or self.step + self.resume_step < self.lr_anneal_steps 169 | ): 170 | batch, cond = next(self.data) 171 | self.run_step(batch, cond, ssim_fn, lpip_fn) 172 | if self.step % self.log_interval == 0: 173 | logger.dumpkvs() 174 | if self.step % self.save_interval == 0: 175 | self.save() 176 | # Run for a finite amount of time in integration tests. 177 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 178 | return 179 | self.step += 1 180 | # Save the last checkpoint if it wasn't already saved. 181 | if (self.step - 1) % self.save_interval != 0: 182 | self.save() 183 | 184 | def run_step(self, batch, cond, ssim_fn=None, lpip_fn=None): 185 | self.forward_backward(batch, cond, ssim_fn, lpip_fn) 186 | if self.use_fp16: 187 | self.optimize_fp16() 188 | else: 189 | self.optimize_normal() 190 | self.log_step() 191 | 192 | def forward_backward(self, batch, cond, ssim_fn=None, lpip_fn=None): 193 | zero_grad(self.model_params) 194 | for i in range(0, batch.shape[0], self.microbatch): 195 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 196 | micro_cond = { 197 | k: v[i : i + self.microbatch].to(dist_util.dev()) 198 | for k, v in cond.items() 199 | } 200 | last_batch = (i + self.microbatch) >= batch.shape[0] 201 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 202 | 203 | compute_losses = functools.partial( 204 | self.diffusion.training_losses, 205 | self.ddp_model, 206 | micro, 207 | t, 208 | model_kwargs=micro_cond, 209 | ssim_fn = ssim_fn, 210 | lpip_fn = lpip_fn 211 | ) 212 | 213 | if last_batch or not self.use_ddp: 214 | losses = compute_losses() 215 | else: 216 | with self.ddp_model.no_sync(): 217 | losses = compute_losses() 218 | 219 | if isinstance(self.schedule_sampler, LossAwareSampler): 220 | self.schedule_sampler.update_with_local_losses( 221 | t, losses["loss"].detach() 222 | ) 223 | 224 | loss = (losses["loss"] * weights).mean() 225 | log_loss_dict( 226 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 227 | ) 228 | if self.use_fp16: 229 | loss_scale = 2 ** self.lg_loss_scale 230 | (loss * loss_scale).backward() 231 | else: 232 | loss.backward() 233 | 234 | def optimize_fp16(self): 235 | if any(not th.isfinite(p.grad).all() for p in self.model_params): 236 | self.lg_loss_scale -= 1 237 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 238 | return 239 | 240 | model_grads_to_master_grads(self.model_params, self.master_params) 241 | self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale)) 242 | self._log_grad_norm() 243 | self._anneal_lr() 244 | self.opt.step() 245 | for rate, params in zip(self.ema_rate, self.ema_params): 246 | update_ema(params, self.master_params, rate=rate) 247 | master_params_to_model_params(self.model_params, self.master_params) 248 | self.lg_loss_scale += self.fp16_scale_growth 249 | 250 | def optimize_normal(self): 251 | self._log_grad_norm() 252 | self._anneal_lr() 253 | self.opt.step() 254 | for rate, params in zip(self.ema_rate, self.ema_params): 255 | update_ema(params, self.master_params, rate=rate) 256 | 257 | def _log_grad_norm(self): 258 | sqsum = 0.0 259 | for p in self.master_params: 260 | sqsum += (p.grad ** 2).sum().item() 261 | logger.logkv_mean("grad_norm", np.sqrt(sqsum)) 262 | 263 | def _anneal_lr(self): 264 | if not self.lr_anneal_steps: 265 | return 266 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 267 | lr = self.lr * (1 - frac_done) 268 | for param_group in self.opt.param_groups: 269 | param_group["lr"] = lr 270 | 271 | def log_step(self): 272 | logger.logkv("step", self.step + self.resume_step) 273 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 274 | if self.use_fp16: 275 | logger.logkv("lg_loss_scale", self.lg_loss_scale) 276 | 277 | def save(self): 278 | def save_checkpoint(rate, params): 279 | state_dict = self._master_params_to_state_dict(params) 280 | if dist.get_rank() == 0: 281 | logger.log(f"saving model {rate}...") 282 | if not rate: 283 | filename = f"model{(self.step+self.resume_step):06d}.pt" 284 | else: 285 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 286 | # with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 287 | with bf.BlobFile(bf.join('/root/autodl-tmp/TASK7/improved_diffusion/model_save_3channels_modi/', filename), "wb") as f: 288 | th.save(state_dict, f) 289 | 290 | save_checkpoint(0, self.master_params) 291 | for rate, params in zip(self.ema_rate, self.ema_params): 292 | save_checkpoint(rate, params) 293 | 294 | if dist.get_rank() == 0: 295 | with bf.BlobFile( 296 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 297 | "wb", 298 | ) as f: 299 | th.save(self.opt.state_dict(), f) 300 | 301 | dist.barrier() 302 | 303 | def _master_params_to_state_dict(self, master_params): 304 | if self.use_fp16: 305 | master_params = unflatten_master_params( 306 | self.model.parameters(), master_params 307 | ) 308 | state_dict = self.model.state_dict() 309 | for i, (name, _value) in enumerate(self.model.named_parameters()): 310 | assert name in state_dict 311 | state_dict[name] = master_params[i] 312 | return state_dict 313 | 314 | def _state_dict_to_master_params(self, state_dict): 315 | params = [state_dict[name] for name, _ in self.model.named_parameters()] 316 | if self.use_fp16: 317 | return make_master_params(params) 318 | else: 319 | return params 320 | 321 | 322 | def parse_resume_step_from_filename(filename): 323 | """ 324 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 325 | checkpoint's number of steps. 326 | """ 327 | split = filename.split("model") 328 | if len(split) < 2: 329 | return 0 330 | split1 = split[-1].split(".")[0] 331 | try: 332 | return int(split1) 333 | except ValueError: 334 | return 0 335 | 336 | 337 | def get_blob_logdir(): 338 | return os.environ.get("DIFFUSION_BLOB_LOGDIR", logger.get_dir()) 339 | 340 | 341 | def find_resume_checkpoint(): 342 | # On your infrastructure, you may want to override this to automatically 343 | # discover the latest checkpoint on your blob storage, etc. 344 | return None 345 | 346 | 347 | def find_ema_checkpoint(main_checkpoint, step, rate): 348 | if main_checkpoint is None: 349 | return None 350 | filename = f"ema_{rate}_{(step):06d}.pt" 351 | path = bf.join(bf.dirname(main_checkpoint), filename) 352 | if bf.exists(path): 353 | return path 354 | return None 355 | 356 | 357 | def log_loss_dict(diffusion, ts, losses): 358 | for key, values in losses.items(): 359 | logger.logkv_mean(key, values.mean().item()) 360 | # Log the quantiles (four quartiles, in particular). 361 | 362 | # for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 363 | # quartile = int(4 * sub_t / diffusion.num_timesteps) 364 | # logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 365 | -------------------------------------------------------------------------------- /IRM/improved_diffusion/unet.py: -------------------------------------------------------------------------------- 1 | from abc import abstractmethod 2 | 3 | import math 4 | 5 | import numpy as np 6 | import torch as th 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | 10 | from .fp16_util import convert_module_to_f16, convert_module_to_f32 11 | from .nn import ( 12 | SiLU, 13 | conv_nd, 14 | linear, 15 | avg_pool_nd, 16 | zero_module, 17 | normalization, 18 | timestep_embedding, 19 | checkpoint, 20 | ) 21 | 22 | 23 | class TimestepBlock(nn.Module): 24 | """ 25 | Any module where forward() takes timestep embeddings as a second argument. 26 | """ 27 | 28 | @abstractmethod 29 | def forward(self, x, emb): 30 | """ 31 | Apply the module to `x` given `emb` timestep embeddings. 32 | """ 33 | 34 | 35 | class TimestepEmbedSequential(nn.Sequential, TimestepBlock): 36 | """ 37 | A sequential module that passes timestep embeddings to the children that 38 | support it as an extra input. 39 | """ 40 | 41 | def forward(self, x, emb): 42 | for layer in self: 43 | if isinstance(layer, TimestepBlock): 44 | x = layer(x, emb) 45 | else: 46 | x = layer(x) 47 | return x 48 | 49 | 50 | class Upsample(nn.Module): 51 | """ 52 | An upsampling layer with an optional convolution. 53 | 54 | :param channels: channels in the inputs and outputs. 55 | :param use_conv: a bool determining if a convolution is applied. 56 | :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then 57 | upsampling occurs in the inner-two dimensions. 58 | """ 59 | 60 | def __init__(self, channels, use_conv, dims=2): 61 | super().__init__() 62 | self.channels = channels 63 | self.use_conv = use_conv 64 | self.dims = dims 65 | if use_conv: 66 | self.conv = conv_nd(dims, channels, channels, 3, padding=1) 67 | 68 | def forward(self, x): 69 | assert x.shape[1] == self.channels 70 | if self.dims == 3: 71 | x = F.interpolate( 72 | x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" 73 | ) 74 | else: 75 | x = F.interpolate(x, scale_factor=2, mode="nearest") 76 | if self.use_conv: 77 | x = self.conv(x) 78 | return x 79 | 80 | 81 | class Downsample(nn.Module): 82 | """ 83 | A downsampling layer with an optional convolution. 84 | 85 | :param channels: channels in the inputs and outputs. 86 | :param use_conv: a bool determining if a convolution is applied. 87 | :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then 88 | downsampling occurs in the inner-two dimensions. 89 | """ 90 | 91 | def __init__(self, channels, use_conv, dims=2): 92 | super().__init__() 93 | self.channels = channels 94 | self.use_conv = use_conv 95 | self.dims = dims 96 | stride = 2 if dims != 3 else (1, 2, 2) 97 | if use_conv: 98 | self.op = conv_nd(dims, channels, channels, 3, stride=stride, padding=1) 99 | else: 100 | self.op = avg_pool_nd(stride) 101 | 102 | def forward(self, x): 103 | assert x.shape[1] == self.channels 104 | return self.op(x) 105 | 106 | 107 | class ResBlock(TimestepBlock): 108 | """ 109 | A residual block that can optionally change the number of channels. 110 | 111 | :param channels: the number of input channels. 112 | :param emb_channels: the number of timestep embedding channels. 113 | :param dropout: the rate of dropout. 114 | :param out_channels: if specified, the number of out channels. 115 | :param use_conv: if True and out_channels is specified, use a spatial 116 | convolution instead of a smaller 1x1 convolution to change the 117 | channels in the skip connection. 118 | :param dims: determines if the signal is 1D, 2D, or 3D. 119 | :param use_checkpoint: if True, use gradient checkpointing on this module. 120 | """ 121 | 122 | def __init__( 123 | self, 124 | channels, 125 | emb_channels, 126 | dropout, 127 | out_channels=None, 128 | use_conv=False, 129 | use_scale_shift_norm=False, 130 | dims=2, 131 | use_checkpoint=False, 132 | ): 133 | super().__init__() 134 | self.channels = channels 135 | self.emb_channels = emb_channels 136 | self.dropout = dropout 137 | self.out_channels = out_channels or channels 138 | self.use_conv = use_conv 139 | self.use_checkpoint = use_checkpoint 140 | self.use_scale_shift_norm = use_scale_shift_norm 141 | 142 | self.in_layers = nn.Sequential( 143 | normalization(channels), 144 | SiLU(), 145 | conv_nd(dims, channels, self.out_channels, 3, padding=1), 146 | ) 147 | self.emb_layers = nn.Sequential( 148 | SiLU(), 149 | linear( 150 | emb_channels, 151 | 2 * self.out_channels if use_scale_shift_norm else self.out_channels, 152 | ), 153 | ) 154 | self.out_layers = nn.Sequential( 155 | normalization(self.out_channels), 156 | SiLU(), 157 | nn.Dropout(p=dropout), 158 | zero_module( 159 | conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) 160 | ), 161 | ) 162 | 163 | if self.out_channels == channels: 164 | self.skip_connection = nn.Identity() 165 | elif use_conv: 166 | self.skip_connection = conv_nd( 167 | dims, channels, self.out_channels, 3, padding=1 168 | ) 169 | else: 170 | self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) 171 | 172 | def forward(self, x, emb): 173 | """ 174 | Apply the block to a Tensor, conditioned on a timestep embedding. 175 | 176 | :param x: an [N x C x ...] Tensor of features. 177 | :param emb: an [N x emb_channels] Tensor of timestep embeddings. 178 | :return: an [N x C x ...] Tensor of outputs. 179 | """ 180 | return checkpoint( 181 | self._forward, (x, emb), self.parameters(), self.use_checkpoint 182 | ) 183 | 184 | def _forward(self, x, emb): 185 | h = self.in_layers(x) 186 | emb_out = self.emb_layers(emb).type(h.dtype) 187 | while len(emb_out.shape) < len(h.shape): 188 | emb_out = emb_out[..., None] 189 | if self.use_scale_shift_norm: 190 | out_norm, out_rest = self.out_layers[0], self.out_layers[1:] 191 | scale, shift = th.chunk(emb_out, 2, dim=1) 192 | h = out_norm(h) * (1 + scale) + shift 193 | h = out_rest(h) 194 | else: 195 | h = h + emb_out 196 | h = self.out_layers(h) 197 | return self.skip_connection(x) + h 198 | 199 | 200 | class AttentionBlock(nn.Module): 201 | """ 202 | An attention block that allows spatial positions to attend to each other. 203 | 204 | Originally ported from here, but adapted to the N-d case. 205 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. 206 | """ 207 | 208 | def __init__(self, channels, num_heads=1, use_checkpoint=False): 209 | super().__init__() 210 | self.channels = channels 211 | self.num_heads = num_heads 212 | self.use_checkpoint = use_checkpoint 213 | 214 | self.norm = normalization(channels) 215 | self.qkv = conv_nd(1, channels, channels * 3, 1) 216 | self.attention = QKVAttention() 217 | self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) 218 | 219 | def forward(self, x): 220 | return checkpoint(self._forward, (x,), self.parameters(), self.use_checkpoint) 221 | 222 | def _forward(self, x): 223 | b, c, *spatial = x.shape 224 | x = x.reshape(b, c, -1) 225 | qkv = self.qkv(self.norm(x)) 226 | qkv = qkv.reshape(b * self.num_heads, -1, qkv.shape[2]) 227 | h = self.attention(qkv) 228 | h = h.reshape(b, -1, h.shape[-1]) 229 | h = self.proj_out(h) 230 | return (x + h).reshape(b, c, *spatial) 231 | 232 | 233 | class QKVAttention(nn.Module): 234 | """ 235 | A module which performs QKV attention. 236 | """ 237 | 238 | def forward(self, qkv): 239 | """ 240 | Apply QKV attention. 241 | 242 | :param qkv: an [N x (C * 3) x T] tensor of Qs, Ks, and Vs. 243 | :return: an [N x C x T] tensor after attention. 244 | """ 245 | ch = qkv.shape[1] // 3 246 | q, k, v = th.split(qkv, ch, dim=1) 247 | scale = 1 / math.sqrt(math.sqrt(ch)) 248 | weight = th.einsum( 249 | "bct,bcs->bts", q * scale, k * scale 250 | ) # More stable with f16 than dividing afterwards 251 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) 252 | return th.einsum("bts,bcs->bct", weight, v) 253 | 254 | @staticmethod 255 | def count_flops(model, _x, y): 256 | """ 257 | A counter for the `thop` package to count the operations in an 258 | attention operation. 259 | 260 | Meant to be used like: 261 | 262 | macs, params = thop.profile( 263 | model, 264 | inputs=(inputs, timestamps), 265 | custom_ops={QKVAttention: QKVAttention.count_flops}, 266 | ) 267 | 268 | """ 269 | b, c, *spatial = y[0].shape 270 | num_spatial = int(np.prod(spatial)) 271 | # We perform two matmuls with the same number of ops. 272 | # The first computes the weight matrix, the second computes 273 | # the combination of the value vectors. 274 | matmul_ops = 2 * b * (num_spatial ** 2) * c 275 | model.total_ops += th.DoubleTensor([matmul_ops]) 276 | 277 | 278 | class UNetModel(nn.Module): 279 | """ 280 | The full UNet model with attention and timestep embedding. 281 | 282 | :param in_channels: channels in the input Tensor. 283 | :param model_channels: base channel count for the model. 284 | :param out_channels: channels in the output Tensor. 285 | :param num_res_blocks: number of residual blocks per downsample. 286 | :param attention_resolutions: a collection of downsample rates at which 287 | attention will take place. May be a set, list, or tuple. 288 | For example, if this contains 4, then at 4x downsampling, attention 289 | will be used. 290 | :param dropout: the dropout probability. 291 | :param channel_mult: channel multiplier for each level of the UNet. 292 | :param conv_resample: if True, use learned convolutions for upsampling and 293 | downsampling. 294 | :param dims: determines if the signal is 1D, 2D, or 3D. 295 | :param num_classes: if specified (as an int), then this model will be 296 | class-conditional with `num_classes` classes. 297 | :param use_checkpoint: use gradient checkpointing to reduce memory usage. 298 | :param num_heads: the number of attention heads in each attention layer. 299 | """ 300 | 301 | def __init__( 302 | self, 303 | in_channels, 304 | model_channels, 305 | out_channels, 306 | num_res_blocks, 307 | attention_resolutions, 308 | dropout=0, 309 | channel_mult=(1, 2, 4, 8), 310 | conv_resample=True, 311 | dims=2, 312 | num_classes=None, 313 | use_checkpoint=False, 314 | num_heads=1, 315 | num_heads_upsample=-1, 316 | use_scale_shift_norm=False, 317 | ): 318 | super().__init__() 319 | 320 | if num_heads_upsample == -1: 321 | num_heads_upsample = num_heads 322 | 323 | self.in_channels = in_channels 324 | self.model_channels = model_channels 325 | self.out_channels = out_channels 326 | self.num_res_blocks = num_res_blocks 327 | self.attention_resolutions = attention_resolutions 328 | self.dropout = dropout 329 | self.channel_mult = channel_mult 330 | self.conv_resample = conv_resample 331 | self.num_classes = num_classes 332 | self.use_checkpoint = use_checkpoint 333 | self.num_heads = num_heads 334 | self.num_heads_upsample = num_heads_upsample 335 | 336 | time_embed_dim = model_channels * 4 337 | self.time_embed = nn.Sequential( 338 | linear(model_channels, time_embed_dim), 339 | SiLU(), 340 | linear(time_embed_dim, time_embed_dim), 341 | ) 342 | 343 | if self.num_classes is not None: 344 | self.label_emb = nn.Embedding(num_classes, time_embed_dim) 345 | 346 | self.input_blocks = nn.ModuleList( 347 | [ 348 | TimestepEmbedSequential( 349 | conv_nd(dims, in_channels, model_channels, 3, padding=1) 350 | ) 351 | ] 352 | ) 353 | input_block_chans = [model_channels] 354 | ch = model_channels 355 | ds = 1 356 | for level, mult in enumerate(channel_mult): 357 | for _ in range(num_res_blocks): 358 | layers = [ 359 | ResBlock( 360 | ch, 361 | time_embed_dim, 362 | dropout, 363 | out_channels=mult * model_channels, 364 | dims=dims, 365 | use_checkpoint=use_checkpoint, 366 | use_scale_shift_norm=use_scale_shift_norm, 367 | ) 368 | ] 369 | ch = mult * model_channels 370 | if ds in attention_resolutions: 371 | layers.append( 372 | AttentionBlock( 373 | ch, use_checkpoint=use_checkpoint, num_heads=num_heads 374 | ) 375 | ) 376 | self.input_blocks.append(TimestepEmbedSequential(*layers)) 377 | input_block_chans.append(ch) 378 | if level != len(channel_mult) - 1: 379 | self.input_blocks.append( 380 | TimestepEmbedSequential(Downsample(ch, conv_resample, dims=dims)) 381 | ) 382 | input_block_chans.append(ch) 383 | ds *= 2 384 | 385 | self.middle_block = TimestepEmbedSequential( 386 | ResBlock( 387 | ch, 388 | time_embed_dim, 389 | dropout, 390 | dims=dims, 391 | use_checkpoint=use_checkpoint, 392 | use_scale_shift_norm=use_scale_shift_norm, 393 | ), 394 | AttentionBlock(ch, use_checkpoint=use_checkpoint, num_heads=num_heads), 395 | ResBlock( 396 | ch, 397 | time_embed_dim, 398 | dropout, 399 | dims=dims, 400 | use_checkpoint=use_checkpoint, 401 | use_scale_shift_norm=use_scale_shift_norm, 402 | ), 403 | ) 404 | 405 | self.output_blocks = nn.ModuleList([]) 406 | for level, mult in list(enumerate(channel_mult))[::-1]: 407 | for i in range(num_res_blocks + 1): 408 | layers = [ 409 | ResBlock( 410 | ch + input_block_chans.pop(), 411 | time_embed_dim, 412 | dropout, 413 | out_channels=model_channels * mult, 414 | dims=dims, 415 | use_checkpoint=use_checkpoint, 416 | use_scale_shift_norm=use_scale_shift_norm, 417 | ) 418 | ] 419 | ch = model_channels * mult 420 | if ds in attention_resolutions: 421 | layers.append( 422 | AttentionBlock( 423 | ch, 424 | use_checkpoint=use_checkpoint, 425 | num_heads=num_heads_upsample, 426 | ) 427 | ) 428 | if level and i == num_res_blocks: 429 | layers.append(Upsample(ch, conv_resample, dims=dims)) 430 | ds //= 2 431 | self.output_blocks.append(TimestepEmbedSequential(*layers)) 432 | 433 | self.out = nn.Sequential( 434 | normalization(ch), 435 | SiLU(), 436 | zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)), 437 | ) 438 | 439 | def convert_to_fp16(self): 440 | """ 441 | Convert the torso of the model to float16. 442 | """ 443 | self.input_blocks.apply(convert_module_to_f16) 444 | self.middle_block.apply(convert_module_to_f16) 445 | self.output_blocks.apply(convert_module_to_f16) 446 | 447 | def convert_to_fp32(self): 448 | """ 449 | Convert the torso of the model to float32. 450 | """ 451 | self.input_blocks.apply(convert_module_to_f32) 452 | self.middle_block.apply(convert_module_to_f32) 453 | self.output_blocks.apply(convert_module_to_f32) 454 | 455 | @property 456 | def inner_dtype(self): 457 | """ 458 | Get the dtype used by the torso of the model. 459 | """ 460 | return next(self.input_blocks.parameters()).dtype 461 | 462 | def forward(self, x, timesteps, y=None): 463 | """ 464 | Apply the model to an input batch. 465 | 466 | :param x: an [N x C x ...] Tensor of inputs. 467 | :param timesteps: a 1-D batch of timesteps. 468 | :param y: an [N] Tensor of labels, if class-conditional. 469 | :return: an [N x C x ...] Tensor of outputs. 470 | """ 471 | assert (y is not None) == ( 472 | self.num_classes is not None 473 | ), "must specify y if and only if the model is class-conditional" 474 | 475 | hs = [] 476 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) 477 | 478 | if self.num_classes is not None: 479 | assert y.shape == (x.shape[0],) 480 | emb = emb + self.label_emb(y) 481 | 482 | h = x.type(self.inner_dtype) 483 | for module in self.input_blocks: 484 | h = module(h, emb) 485 | hs.append(h) 486 | h = self.middle_block(h, emb) 487 | for module in self.output_blocks: 488 | cat_in = th.cat([h, hs.pop()], dim=1) 489 | h = module(cat_in, emb) 490 | h = h.type(x.dtype) 491 | return self.out(h) 492 | 493 | def get_feature_vectors(self, x, timesteps, y=None): 494 | """ 495 | Apply the model and return all of the intermediate tensors. 496 | 497 | :param x: an [N x C x ...] Tensor of inputs. 498 | :param timesteps: a 1-D batch of timesteps. 499 | :param y: an [N] Tensor of labels, if class-conditional. 500 | :return: a dict with the following keys: 501 | - 'down': a list of hidden state tensors from downsampling. 502 | - 'middle': the tensor of the output of the lowest-resolution 503 | block in the model. 504 | - 'up': a list of hidden state tensors from upsampling. 505 | """ 506 | hs = [] 507 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) 508 | if self.num_classes is not None: 509 | assert y.shape == (x.shape[0],) 510 | emb = emb + self.label_emb(y) 511 | result = dict(down=[], up=[]) 512 | h = x.type(self.inner_dtype) 513 | for module in self.input_blocks: 514 | h = module(h, emb) 515 | hs.append(h) 516 | result["down"].append(h.type(x.dtype)) 517 | h = self.middle_block(h, emb) 518 | result["middle"] = h.type(x.dtype) 519 | for module in self.output_blocks: 520 | cat_in = th.cat([h, hs.pop()], dim=1) 521 | h = module(cat_in, emb) 522 | result["up"].append(h.type(x.dtype)) 523 | return result 524 | 525 | 526 | class SuperResModel(UNetModel): 527 | """ 528 | A UNetModel that performs super-resolution. 529 | 530 | Expects an extra kwarg `low_res` to condition on a low-resolution image. 531 | """ 532 | 533 | def __init__(self, in_channels, *args, **kwargs): 534 | super().__init__(in_channels * 3, *args, **kwargs) #############注意注意 535 | 536 | def forward(self, x, timesteps, low_res=None, _36_res=None, **kwargs): 537 | _, _, new_height, new_width = x.shape 538 | upsampled = low_res 539 | x = th.cat([x, upsampled, _36_res], dim=1) 540 | return super().forward(x, timesteps, **kwargs) + upsampled # _72_res 541 | 542 | # def get_feature_vectors(self, x, timesteps, low_res=None, _72_res=None, **kwargs): 543 | # _, new_height, new_width, _ = x.shape 544 | # upsampled = low_res 545 | # x = th.cat([x, upsampled - _72_res, _72_res], dim=1) 546 | # return super().get_feature_vectors(x, timesteps, **kwargs) 547 | 548 | -------------------------------------------------------------------------------- /IRM/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="improved-diffusion", 5 | py_modules=["improved_diffusion"], 6 | install_requires=["blobfile>=1.0.5", "torch", "tqdm"], 7 | ) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Code for DDDM 2 | 3 | This is the code for paper *A Dual-domain Diffusion Model for Sparse-view CT Reconstruction*. 4 | 5 | The pre-processing code and MATLAB fan2parallel code will be released later. 6 | 7 | ## Installation SUM 8 | 9 | Clone this repository. 10 | 11 | Navigate to SUM folder in your terminal. 12 | 13 | Then run: 14 | 15 | ``` 16 | pip install -e . 17 | ``` 18 | 19 | This should install the `improved_diffusion` python package that the scripts depend on. 20 | 21 | Note that some packages need to be installed manually. 22 | 23 | Note that the SUM and IRM are now trained and inferenced **separately**. Thus, each time you change the working module, please run `pip install -e .` **again** under corresponding file path. 24 | 25 | 26 | 27 | ## Preparing Data for SUM 28 | 29 | Regardless of the type of CT data, appropriate preprocessing is required. SUM requires specially "normalized" parallel-beam sinograms as input (Given the high dynamic range of CT data, we set the "black" area of the sinogram to -1, and let the mean of all sinograms for the same organ to linearly vary to 0, allowing the maximum to drift. Please see xxx for details). Also, you can use fan-beam sinograms as input, with appropriate change in `reco_train.py` to meet the demand of FBP. 30 | 31 | 32 | 33 | ## Training SUM 34 | 35 | The default hyperparameters are appropriately set with adequate testing. However, some hyperparameters related to U-Net may need to be changed in line with your dataset and task. 36 | 37 | 38 | 39 | ## Sampling SUM 40 | 41 | The above training script saves checkpoints to `.pt` files in the logging directory. These checkpoints will have names like `ema_0.9999_200000.pt` and `model200000.pt`. You will likely want to sample from the EMA models, since those produce much better samples. 42 | 43 | Once you have a path to your model, you can generate a single batch of sample like so: 44 | 45 | ``` 46 | python SUM/scripts/super_res_sample.py 47 | ``` 48 | 49 | Or generate a large batch of samples like so: 50 | 51 | ``` 52 | python SUM/scripts/multi_super_res.py 53 | ``` 54 | 55 | Note that for SUM, the sampling strategy is exclusive. 56 | 57 | 58 | 59 | ## Installation IRM 60 | 61 | Clone this repository. 62 | 63 | Navigate to IRM folder in your terminal. 64 | 65 | Then run: 66 | 67 | ``` 68 | pip install -e . 69 | ``` 70 | 71 | This should install the `improved_diffusion` python package that the scripts depend on. 72 | 73 | Note that some packages need to be installed manually. 74 | 75 | Note that the SUM and IRM are now trained and inferenced separately. Thus, each time you change the working module, please run `pip install -e .` again under corresponding file path. 76 | 77 | 78 | 79 | ## Preparing Data for IRM 80 | 81 | Once the training process of SUM has completed, all the training & test & validation dataset should be sampled, transferred by FBP as the input of IRM. 82 | 83 | Note that the CT image data should also be “normalized” followed by the identical strategy as that in SUM. 84 | 85 | 86 | 87 | ## Training IRM 88 | 89 | The default hyperparameters are appropriately set with adequate testing. However, some hyperparameters related to U-Net may need to be changed in line with your dataset and task. 90 | 91 | 92 | 93 | ## Sampling IRM 94 | 95 | The above training script saves checkpoints to `.pt` files in the logging directory. These checkpoints will have names like `ema_0.9999_200000.pt` and `model200000.pt`. You will likely want to sample from the EMA models, since those produce much better samples. 96 | 97 | Once you have a path to your model, you can generate a single batch of sample like so: 98 | 99 | ``` 100 | python IRM/scripts/img_super_res_sample.py 101 | ``` 102 | 103 | Or generate a large batch of samples like so: 104 | 105 | ``` 106 | python IRM/scripts/img_multi_sample.py 107 | ``` 108 | 109 | Note that for IRM, strides IDDPM and DDIM are optional. We recommend to use DDIM and set parameter `--timestep_respacing` to 2 as more timesteps may lead to unexpected faked details. 110 | -------------------------------------------------------------------------------- /SUM/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /SUM/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = 8 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | 28 | comm = MPI.COMM_WORLD 29 | backend = "gloo" if not th.cuda.is_available() else "nccl" 30 | 31 | if backend == "gloo": 32 | hostname = "localhost" 33 | else: 34 | hostname = socket.gethostbyname(socket.getfqdn()) 35 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 36 | os.environ["RANK"] = str(comm.rank) 37 | os.environ["WORLD_SIZE"] = str(comm.size) 38 | 39 | port = comm.bcast(_find_free_port(), root=0) 40 | os.environ["MASTER_PORT"] = str(port) 41 | dist.init_process_group(backend=backend, init_method="env://") 42 | 43 | 44 | def dev(): 45 | """ 46 | Get the device to use for torch.distributed. 47 | """ 48 | if th.cuda.is_available(): 49 | return th.device(f"cuda:{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}") 50 | return th.device("cpu") 51 | 52 | 53 | def load_state_dict(path, **kwargs): 54 | """ 55 | Load a PyTorch file without redundant fetches across MPI ranks. 56 | """ 57 | if MPI.COMM_WORLD.Get_rank() == 0: 58 | with bf.BlobFile(path, "rb") as f: 59 | data = f.read() 60 | else: 61 | data = None 62 | data = MPI.COMM_WORLD.bcast(data) 63 | return th.load(io.BytesIO(data), **kwargs) 64 | 65 | 66 | def sync_params(params): 67 | """ 68 | Synchronize a sequence of Tensors across ranks from rank 0. 69 | """ 70 | for p in params: 71 | with th.no_grad(): 72 | dist.broadcast(p, 0) 73 | 74 | 75 | def _find_free_port(): 76 | try: 77 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 78 | s.bind(("", 0)) 79 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 80 | return s.getsockname()[1] 81 | finally: 82 | s.close() 83 | -------------------------------------------------------------------------------- /SUM/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import torch.nn as nn 6 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 7 | 8 | 9 | def convert_module_to_f16(l): 10 | """ 11 | Convert primitive modules to float16. 12 | """ 13 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 14 | l.weight.data = l.weight.data.half() 15 | l.bias.data = l.bias.data.half() 16 | 17 | 18 | def convert_module_to_f32(l): 19 | """ 20 | Convert primitive modules to float32, undoing convert_module_to_f16(). 21 | """ 22 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 23 | l.weight.data = l.weight.data.float() 24 | l.bias.data = l.bias.data.float() 25 | 26 | 27 | def make_master_params(model_params): 28 | """ 29 | Copy model parameters into a (differently-shaped) list of full-precision 30 | parameters. 31 | """ 32 | master_params = _flatten_dense_tensors( 33 | [param.detach().float() for param in model_params] 34 | ) 35 | master_params = nn.Parameter(master_params) 36 | master_params.requires_grad = True 37 | return [master_params] 38 | 39 | 40 | def model_grads_to_master_grads(model_params, master_params): 41 | """ 42 | Copy the gradients from the model parameters into the master parameters 43 | from make_master_params(). 44 | """ 45 | master_params[0].grad = _flatten_dense_tensors( 46 | [param.grad.data.detach().float() for param in model_params] 47 | ) 48 | 49 | 50 | def master_params_to_model_params(model_params, master_params): 51 | """ 52 | Copy the master parameter data back into the model parameters. 53 | """ 54 | # Without copying to a list, if a generator is passed, this will 55 | # silently not copy any parameters. 56 | model_params = list(model_params) 57 | 58 | for param, master_param in zip( 59 | model_params, unflatten_master_params(model_params, master_params) 60 | ): 61 | param.detach().copy_(master_param) 62 | 63 | 64 | def unflatten_master_params(model_params, master_params): 65 | """ 66 | Unflatten the master parameters to look like model_params. 67 | """ 68 | return _unflatten_dense_tensors(master_params[0].detach(), model_params) 69 | 70 | 71 | def zero_grad(model_params): 72 | for param in model_params: 73 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 74 | if param.grad is not None: 75 | param.grad.detach_() 76 | param.grad.zero_() 77 | -------------------------------------------------------------------------------- /SUM/helper.py: -------------------------------------------------------------------------------- 1 | from skimage.io import imread, imsave 2 | import numpy as np 3 | import pandas as pd 4 | import json 5 | import tifffile 6 | 7 | 8 | def load_tiff_stack(file): 9 | ''' 10 | 11 | :param file: Path object describing the location of the file 12 | :return: a numpy array of the volume 13 | ''' 14 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 15 | raise FileNotFoundError('Input has to be tif.') 16 | else: 17 | volume = imread(file, plugin='tifffile') 18 | return volume 19 | 20 | 21 | def load_tiff_stack_with_metadata(file): 22 | ''' 23 | 24 | :param file: Path object describing the location of the file 25 | :return: a numpy array of the volume, a dict with the metadata 26 | ''' 27 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 28 | raise FileNotFoundError('File has to be tif.') 29 | with tifffile.TiffFile(file) as tif: 30 | data = tif.asarray() 31 | metadata = tif.pages[0].tags["ImageDescription"].value 32 | metadata = metadata.replace("'", "\"") 33 | try: 34 | metadata = json.loads(metadata) 35 | except: 36 | print('The tiff file you try to open does not seem to have metadata attached.') 37 | metadata = None 38 | return data, metadata 39 | 40 | 41 | def save_to_tiff_stack(array, file): 42 | ''' 43 | 44 | :param array: Array to save 45 | :param file: Path object to save to 46 | ''' 47 | if not file.parent.is_dir(): 48 | file.parent.mkdir(parents=True, exist_ok=True) 49 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 50 | raise FileNotFoundError('File has to be tif.') 51 | else: 52 | imsave(file, array, plugin='tifffile', check_contrast=False) 53 | 54 | 55 | def save_to_tiff_stack_with_metadata(array, file, metadata): 56 | ''' 57 | 58 | :param array: 59 | :param file: 60 | :return: 61 | ''' 62 | if not file.parent.is_dir(): 63 | file.parent.mkdir(parents=True, exist_ok=True) 64 | if not (file.name.endswith('.tif') or file.name.endswith('.tiff')): 65 | raise FileNotFoundError('File has to be tif.') 66 | else: 67 | # metadata = json.dumps(metadata) 68 | # tifffile.imsave(file, array, description=metadata) 69 | tifffile.imwrite(file, shape=array.shape, dtype=array.dtype, metadata=metadata) 70 | 71 | # memory map numpy array to data in OME-TIFF file 72 | memmap_stack = tifffile.memmap(file) 73 | 74 | # write data to memory-mapped array 75 | for t in range(array.shape[0]): 76 | memmap_stack[t] = array[t, :, :] 77 | memmap_stack.flush() 78 | 79 | 80 | def load_csv(file): 81 | ''' 82 | 83 | :param file: Path object describing the location of the file 84 | :return: pandas dataframe containing the information from the file 85 | ''' 86 | if not file.name.endswith('.csv'): 87 | raise FileNotFoundError('Input has to be a csv file.') 88 | else: 89 | data = pd.read_csv(file) 90 | return data 91 | 92 | 93 | def save_to_json(params, file): 94 | ''' 95 | 96 | :param params: parameter dict 97 | :param file: path to a file to save to 98 | :return: 99 | ''' 100 | if not file.parent.is_dir(): 101 | file.parent.mkdir(parents=True, exist_ok=True) 102 | with open(file, 'w') as f: 103 | json.dump(params, f, indent=2) 104 | 105 | 106 | def load_from_json(file): 107 | ''' 108 | 109 | :param file: path to a file to load 110 | :return: contents of the file 111 | ''' 112 | with open(file) as f: 113 | data = json.load(f) 114 | return data 115 | -------------------------------------------------------------------------------- /SUM/image_datasets.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import blobfile as bf 3 | from mpi4py import MPI 4 | import numpy as np 5 | from torch.utils.data import DataLoader, Dataset 6 | import torch 7 | import os 8 | import random 9 | import torch.nn.functional as F 10 | 11 | 12 | 13 | def improved_data_preprocess(sino_path = '/root/autodl-tmp/TASK4/SINO_DATA/train_stage1', batch_size = 1, shuffle = True, num_workers = 8): 14 | train_dataset = diffCT_Dataset(sino_path) 15 | train_data_loader = DataLoader(train_dataset, batch_size, shuffle, num_workers=num_workers, pin_memory=True, drop_last=True) 16 | while True: 17 | yield from train_data_loader 18 | 19 | def count_npy_files(folder_path: str): 20 | files = os.listdir(folder_path) 21 | npy_files = [file for file in files if file.endswith('.npy')] 22 | return len(npy_files) 23 | 24 | class diffCT_Dataset(Dataset): 25 | def __init__(self, sino_path): 26 | self.sino_path = sino_path 27 | self.sino_num = count_npy_files(sino_path) 28 | 29 | def __len__(self): 30 | return (3276+3276+3255+4154) #self.sino_num 31 | def __getitem__(self,index): 32 | out_dict = {} 33 | 34 | if index < (3276+3276) : 35 | x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_C_parallel_sino/{index}.npy')).to(torch.float32),dim=0) - 1. 36 | # x_nd = torch.tensor(np.load(self.sino_path + f'/full_C_fan_sino/{index}.npy')).to(torch.float32) 37 | # out_dict["y"] = np.array(0, dtype=np.int64) 38 | 39 | elif index < (3255+4154): 40 | x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index-3276}.npy')).to(torch.float32),dim=0) - 1. 41 | # x_nd = torch.tensor(np.load(self.sino_path + f'/full_L_fan_sino/{index-3276}.npy')).to(torch.float32) 42 | # out_dict["y"] = np.array(1, dtype=np.int64) 43 | 44 | return x_nd, out_dict 45 | 46 | 47 | 48 | 49 | 50 | def improved_data_preprocess_val(sino_path = '/root/autodl-tmp/TASK4/SINO_DATA/train_stage1', batch_size = 1, shuffle = True, num_workers = 8): 51 | train_dataset = diffCT_Dataset_val(sino_path) 52 | train_data_loader = DataLoader(train_dataset, batch_size, shuffle, num_workers=num_workers, pin_memory=True, drop_last=True) 53 | while True: 54 | yield from train_data_loader 55 | 56 | class diffCT_Dataset_val(Dataset): 57 | def __init__(self, sino_path): 58 | self.sino_path = sino_path 59 | self.sino_num = count_npy_files(sino_path) 60 | def __len__(self): 61 | return 1000 #self.sino_num 62 | def __getitem__(self,index): 63 | out_dict = {} 64 | 65 | choice1 = random.random() 66 | index = 100 67 | if choice1 < (1/3): 68 | x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_C_parallel_sino/{index}.npy')).to(torch.float32),dim=0) - 1. 69 | # out_dict["y"] = np.array(0, dtype=np.int64) 70 | 71 | elif choice1 < (2/3): 72 | x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index}.npy')).to(torch.float32),dim=0) - 1. 73 | # out_dict["y"] = np.array(1, dtype=np.int64) 74 | 75 | else: 76 | x_nd = torch.tensor(np.load(f'/root/autodl-tmp/TASK4/SINO_DATA/val_stage2/full_L/val_{index}.npy')).to(torch.float32) 77 | # out_dict["y"] = np.array(2, dtype=np.int64) 78 | 79 | return x_nd, out_dict 80 | 81 | 82 | 83 | 84 | def trans_data_preprocess(sino_path = '/root/autodl-tmp/TASK4/SINO_DATA/val_stage2', batch_size = 1, shuffle = False, num_workers = 8): 85 | train_dataset = trans_diffCT_Dataset(sino_path) 86 | train_data_loader = DataLoader(train_dataset, batch_size, shuffle, num_workers=num_workers, pin_memory=True, drop_last=True) 87 | while True: 88 | yield from train_data_loader 89 | 90 | 91 | class trans_diffCT_Dataset(Dataset): 92 | def __init__(self, sino_path): 93 | self.sino_path = sino_path 94 | # self.sino_num = count_npy_files(sino_path) 95 | def __len__(self): 96 | return 800 #-496 97 | def __getitem__(self,index): 98 | out_dict = {} 99 | index = index + 0 100 | x_nd = torch.unsqueeze(torch.tensor(np.load(self.sino_path + f'/full_L_parallel_sino/{index}.npy')).to(torch.float32),dim=0) - 1. 101 | 102 | # out_dict["y"] = np.array(1, dtype=np.int64) 103 | out_dict["index"] = index 104 | 105 | return x_nd, out_dict 106 | -------------------------------------------------------------------------------- /SUM/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Logger copied from OpenAI baselines to avoid extra RL-based dependencies: 3 | https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py 4 | """ 5 | 6 | import os 7 | import sys 8 | import shutil 9 | import os.path as osp 10 | import json 11 | import time 12 | import datetime 13 | import tempfile 14 | import warnings 15 | from collections import defaultdict 16 | from contextlib import contextmanager 17 | 18 | DEBUG = 10 19 | INFO = 20 20 | WARN = 30 21 | ERROR = 40 22 | 23 | DISABLED = 50 24 | 25 | 26 | class KVWriter(object): 27 | def writekvs(self, kvs): 28 | raise NotImplementedError 29 | 30 | 31 | class SeqWriter(object): 32 | def writeseq(self, seq): 33 | raise NotImplementedError 34 | 35 | 36 | class HumanOutputFormat(KVWriter, SeqWriter): 37 | def __init__(self, filename_or_file): 38 | if isinstance(filename_or_file, str): 39 | self.file = open(filename_or_file, "wt") 40 | self.own_file = True 41 | else: 42 | assert hasattr(filename_or_file, "read"), ( 43 | "expected file or str, got %s" % filename_or_file 44 | ) 45 | self.file = filename_or_file 46 | self.own_file = False 47 | 48 | def writekvs(self, kvs): 49 | # Create strings for printing 50 | key2str = {} 51 | for (key, val) in sorted(kvs.items()): 52 | if hasattr(val, "__float__"): 53 | valstr = "%-8.3g" % val 54 | else: 55 | valstr = str(val) 56 | key2str[self._truncate(key)] = self._truncate(valstr) 57 | 58 | # Find max widths 59 | if len(key2str) == 0: 60 | print("WARNING: tried to write empty key-value dict") 61 | return 62 | else: 63 | keywidth = max(map(len, key2str.keys())) 64 | valwidth = max(map(len, key2str.values())) 65 | 66 | # Write out the data 67 | dashes = "-" * (keywidth + valwidth + 7) 68 | lines = [dashes] 69 | for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): 70 | lines.append( 71 | "| %s%s | %s%s |" 72 | % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) 73 | ) 74 | lines.append(dashes) 75 | self.file.write("\n".join(lines) + "\n") 76 | 77 | # Flush the output to the file 78 | self.file.flush() 79 | 80 | def _truncate(self, s): 81 | maxlen = 30 82 | return s[: maxlen - 3] + "..." if len(s) > maxlen else s 83 | 84 | def writeseq(self, seq): 85 | seq = list(seq) 86 | for (i, elem) in enumerate(seq): 87 | self.file.write(elem) 88 | if i < len(seq) - 1: # add space unless this is the last one 89 | self.file.write(" ") 90 | self.file.write("\n") 91 | self.file.flush() 92 | 93 | def close(self): 94 | if self.own_file: 95 | self.file.close() 96 | 97 | 98 | class JSONOutputFormat(KVWriter): 99 | def __init__(self, filename): 100 | self.file = open(filename, "wt") 101 | 102 | def writekvs(self, kvs): 103 | for k, v in sorted(kvs.items()): 104 | if hasattr(v, "dtype"): 105 | kvs[k] = float(v) 106 | self.file.write(json.dumps(kvs) + "\n") 107 | self.file.flush() 108 | 109 | def close(self): 110 | self.file.close() 111 | 112 | 113 | class CSVOutputFormat(KVWriter): 114 | def __init__(self, filename): 115 | self.file = open(filename, "w+t") 116 | self.keys = [] 117 | self.sep = "," 118 | 119 | def writekvs(self, kvs): 120 | # Add our current row to the history 121 | extra_keys = list(kvs.keys() - self.keys) 122 | extra_keys.sort() 123 | if extra_keys: 124 | self.keys.extend(extra_keys) 125 | self.file.seek(0) 126 | lines = self.file.readlines() 127 | self.file.seek(0) 128 | for (i, k) in enumerate(self.keys): 129 | if i > 0: 130 | self.file.write(",") 131 | self.file.write(k) 132 | self.file.write("\n") 133 | for line in lines[1:]: 134 | self.file.write(line[:-1]) 135 | self.file.write(self.sep * len(extra_keys)) 136 | self.file.write("\n") 137 | for (i, k) in enumerate(self.keys): 138 | if i > 0: 139 | self.file.write(",") 140 | v = kvs.get(k) 141 | if v is not None: 142 | self.file.write(str(v)) 143 | self.file.write("\n") 144 | self.file.flush() 145 | 146 | def close(self): 147 | self.file.close() 148 | 149 | 150 | class TensorBoardOutputFormat(KVWriter): 151 | """ 152 | Dumps key/value pairs into TensorBoard's numeric format. 153 | """ 154 | 155 | def __init__(self, dir): 156 | os.makedirs(dir, exist_ok=True) 157 | self.dir = dir 158 | self.step = 1 159 | prefix = "events" 160 | path = osp.join(osp.abspath(dir), prefix) 161 | import tensorflow as tf 162 | from tensorflow.python import pywrap_tensorflow 163 | from tensorflow.core.util import event_pb2 164 | from tensorflow.python.util import compat 165 | 166 | self.tf = tf 167 | self.event_pb2 = event_pb2 168 | self.pywrap_tensorflow = pywrap_tensorflow 169 | self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) 170 | 171 | def writekvs(self, kvs): 172 | def summary_val(k, v): 173 | kwargs = {"tag": k, "simple_value": float(v)} 174 | return self.tf.Summary.Value(**kwargs) 175 | 176 | summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) 177 | event = self.event_pb2.Event(wall_time=time.time(), summary=summary) 178 | event.step = ( 179 | self.step 180 | ) # is there any reason why you'd want to specify the step? 181 | self.writer.WriteEvent(event) 182 | self.writer.Flush() 183 | self.step += 1 184 | 185 | def close(self): 186 | if self.writer: 187 | self.writer.Close() 188 | self.writer = None 189 | 190 | 191 | def make_output_format(format, ev_dir, log_suffix=""): 192 | os.makedirs(ev_dir, exist_ok=True) 193 | if format == "stdout": 194 | return HumanOutputFormat(sys.stdout) 195 | elif format == "log": 196 | return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) 197 | elif format == "json": 198 | return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) 199 | elif format == "csv": 200 | return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) 201 | elif format == "tensorboard": 202 | return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) 203 | else: 204 | raise ValueError("Unknown format specified: %s" % (format,)) 205 | 206 | 207 | # ================================================================ 208 | # API 209 | # ================================================================ 210 | 211 | 212 | def logkv(key, val): 213 | """ 214 | Log a value of some diagnostic 215 | Call this once for each diagnostic quantity, each iteration 216 | If called many times, last value will be used. 217 | """ 218 | get_current().logkv(key, val) 219 | 220 | 221 | def logkv_mean(key, val): 222 | """ 223 | The same as logkv(), but if called many times, values averaged. 224 | """ 225 | get_current().logkv_mean(key, val) 226 | 227 | 228 | def logkvs(d): 229 | """ 230 | Log a dictionary of key-value pairs 231 | """ 232 | for (k, v) in d.items(): 233 | logkv(k, v) 234 | 235 | 236 | def dumpkvs(): 237 | """ 238 | Write all of the diagnostics from the current iteration 239 | """ 240 | return get_current().dumpkvs() 241 | 242 | 243 | def getkvs(): 244 | return get_current().name2val 245 | 246 | 247 | def log(*args, level=INFO): 248 | """ 249 | Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). 250 | """ 251 | get_current().log(*args, level=level) 252 | 253 | 254 | def debug(*args): 255 | log(*args, level=DEBUG) 256 | 257 | 258 | def info(*args): 259 | log(*args, level=INFO) 260 | 261 | 262 | def warn(*args): 263 | log(*args, level=WARN) 264 | 265 | 266 | def error(*args): 267 | log(*args, level=ERROR) 268 | 269 | 270 | def set_level(level): 271 | """ 272 | Set logging threshold on current logger. 273 | """ 274 | get_current().set_level(level) 275 | 276 | 277 | def set_comm(comm): 278 | get_current().set_comm(comm) 279 | 280 | 281 | def get_dir(): 282 | """ 283 | Get directory that log files are being written to. 284 | will be None if there is no output directory (i.e., if you didn't call start) 285 | """ 286 | return get_current().get_dir() 287 | 288 | 289 | record_tabular = logkv 290 | dump_tabular = dumpkvs 291 | 292 | 293 | @contextmanager 294 | def profile_kv(scopename): 295 | logkey = "wait_" + scopename 296 | tstart = time.time() 297 | try: 298 | yield 299 | finally: 300 | get_current().name2val[logkey] += time.time() - tstart 301 | 302 | 303 | def profile(n): 304 | """ 305 | Usage: 306 | @profile("my_func") 307 | def my_func(): code 308 | """ 309 | 310 | def decorator_with_name(func): 311 | def func_wrapper(*args, **kwargs): 312 | with profile_kv(n): 313 | return func(*args, **kwargs) 314 | 315 | return func_wrapper 316 | 317 | return decorator_with_name 318 | 319 | 320 | # ================================================================ 321 | # Backend 322 | # ================================================================ 323 | 324 | 325 | def get_current(): 326 | if Logger.CURRENT is None: 327 | _configure_default_logger() 328 | 329 | return Logger.CURRENT 330 | 331 | 332 | class Logger(object): 333 | DEFAULT = None # A logger with no output files. (See right below class definition) 334 | # So that you can still log to the terminal without setting up any output files 335 | CURRENT = None # Current logger being used by the free functions above 336 | 337 | def __init__(self, dir, output_formats, comm=None): 338 | self.name2val = defaultdict(float) # values this iteration 339 | self.name2cnt = defaultdict(int) 340 | self.level = INFO 341 | self.dir = dir 342 | self.output_formats = output_formats 343 | self.comm = comm 344 | 345 | # Logging API, forwarded 346 | # ---------------------------------------- 347 | def logkv(self, key, val): 348 | self.name2val[key] = val 349 | 350 | def logkv_mean(self, key, val): 351 | oldval, cnt = self.name2val[key], self.name2cnt[key] 352 | self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) 353 | self.name2cnt[key] = cnt + 1 354 | 355 | def dumpkvs(self): 356 | if self.comm is None: 357 | d = self.name2val 358 | else: 359 | d = mpi_weighted_mean( 360 | self.comm, 361 | { 362 | name: (val, self.name2cnt.get(name, 1)) 363 | for (name, val) in self.name2val.items() 364 | }, 365 | ) 366 | if self.comm.rank != 0: 367 | d["dummy"] = 1 # so we don't get a warning about empty dict 368 | out = d.copy() # Return the dict for unit testing purposes 369 | for fmt in self.output_formats: 370 | if isinstance(fmt, KVWriter): 371 | fmt.writekvs(d) 372 | self.name2val.clear() 373 | self.name2cnt.clear() 374 | return out 375 | 376 | def log(self, *args, level=INFO): 377 | if self.level <= level: 378 | self._do_log(args) 379 | 380 | # Configuration 381 | # ---------------------------------------- 382 | def set_level(self, level): 383 | self.level = level 384 | 385 | def set_comm(self, comm): 386 | self.comm = comm 387 | 388 | def get_dir(self): 389 | return self.dir 390 | 391 | def close(self): 392 | for fmt in self.output_formats: 393 | fmt.close() 394 | 395 | # Misc 396 | # ---------------------------------------- 397 | def _do_log(self, args): 398 | for fmt in self.output_formats: 399 | if isinstance(fmt, SeqWriter): 400 | fmt.writeseq(map(str, args)) 401 | 402 | 403 | def get_rank_without_mpi_import(): 404 | # check environment variables here instead of importing mpi4py 405 | # to avoid calling MPI_Init() when this module is imported 406 | for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: 407 | if varname in os.environ: 408 | return int(os.environ[varname]) 409 | return 0 410 | 411 | 412 | def mpi_weighted_mean(comm, local_name2valcount): 413 | """ 414 | Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 415 | Perform a weighted average over dicts that are each on a different node 416 | Input: local_name2valcount: dict mapping key -> (value, count) 417 | Returns: key -> mean 418 | """ 419 | all_name2valcount = comm.gather(local_name2valcount) 420 | if comm.rank == 0: 421 | name2sum = defaultdict(float) 422 | name2count = defaultdict(float) 423 | for n2vc in all_name2valcount: 424 | for (name, (val, count)) in n2vc.items(): 425 | try: 426 | val = float(val) 427 | except ValueError: 428 | if comm.rank == 0: 429 | warnings.warn( 430 | "WARNING: tried to compute mean on non-float {}={}".format( 431 | name, val 432 | ) 433 | ) 434 | else: 435 | name2sum[name] += val * count 436 | name2count[name] += count 437 | return {name: name2sum[name] / name2count[name] for name in name2sum} 438 | else: 439 | return {} 440 | 441 | 442 | def configure(dir=None, format_strs=None, comm=None, log_suffix=""): 443 | """ 444 | If comm is provided, average all numerical stats across that comm 445 | """ 446 | if dir is None: 447 | dir = os.getenv("OPENAI_LOGDIR") 448 | if dir is None: 449 | dir = osp.join( 450 | tempfile.gettempdir(), 451 | datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 452 | ) 453 | assert isinstance(dir, str) 454 | dir = os.path.expanduser(dir) 455 | os.makedirs(os.path.expanduser(dir), exist_ok=True) 456 | 457 | rank = get_rank_without_mpi_import() 458 | if rank > 0: 459 | log_suffix = log_suffix + "-rank%03i" % rank 460 | 461 | if format_strs is None: 462 | if rank == 0: 463 | format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") 464 | else: 465 | format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") 466 | format_strs = filter(None, format_strs) 467 | output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] 468 | 469 | Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) 470 | if output_formats: 471 | log("Logging to %s" % dir) 472 | 473 | 474 | def _configure_default_logger(): 475 | configure() 476 | Logger.DEFAULT = Logger.CURRENT 477 | 478 | 479 | def reset(): 480 | if Logger.CURRENT is not Logger.DEFAULT: 481 | Logger.CURRENT.close() 482 | Logger.CURRENT = Logger.DEFAULT 483 | log("Reset logger") 484 | 485 | 486 | @contextmanager 487 | def scoped_configure(dir=None, format_strs=None, comm=None): 488 | prevlogger = Logger.CURRENT 489 | configure(dir=dir, format_strs=format_strs, comm=comm) 490 | try: 491 | yield 492 | finally: 493 | Logger.CURRENT.close() 494 | Logger.CURRENT = prevlogger 495 | 496 | -------------------------------------------------------------------------------- /SUM/losses.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for various likelihood-based losses. These are ported from the original 3 | Ho et al. diffusion models codebase: 4 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 5 | """ 6 | 7 | import numpy as np 8 | 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /SUM/nn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various utilities for neural networks. 3 | """ 4 | 5 | import math 6 | 7 | import torch as th 8 | import torch.nn as nn 9 | 10 | 11 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 12 | class SiLU(nn.Module): 13 | def forward(self, x): 14 | return x * th.sigmoid(x) 15 | 16 | 17 | class GroupNorm32(nn.GroupNorm): 18 | def forward(self, x): 19 | return super().forward(x.float()).type(x.dtype) 20 | 21 | 22 | def conv_nd(dims, *args, **kwargs): 23 | """ 24 | Create a 1D, 2D, or 3D convolution module. 25 | """ 26 | if dims == 1: 27 | return nn.Conv1d(*args, **kwargs) 28 | elif dims == 2: 29 | return nn.Conv2d(*args, **kwargs) 30 | elif dims == 3: 31 | return nn.Conv3d(*args, **kwargs) 32 | raise ValueError(f"unsupported dimensions: {dims}") 33 | 34 | 35 | def linear(*args, **kwargs): 36 | """ 37 | Create a linear module. 38 | """ 39 | return nn.Linear(*args, **kwargs) 40 | 41 | 42 | def avg_pool_nd(dims, *args, **kwargs): 43 | """ 44 | Create a 1D, 2D, or 3D average pooling module. 45 | """ 46 | if dims == 1: 47 | return nn.AvgPool1d(*args, **kwargs) 48 | elif dims == 2: 49 | return nn.AvgPool2d(*args, **kwargs) 50 | elif dims == 3: 51 | return nn.AvgPool3d(*args, **kwargs) 52 | raise ValueError(f"unsupported dimensions: {dims}") 53 | 54 | 55 | def update_ema(target_params, source_params, rate=0.99): 56 | """ 57 | Update target parameters to be closer to those of source parameters using 58 | an exponential moving average. 59 | 60 | :param target_params: the target parameter sequence. 61 | :param source_params: the source parameter sequence. 62 | :param rate: the EMA rate (closer to 1 means slower). 63 | """ 64 | for targ, src in zip(target_params, source_params): 65 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 66 | 67 | 68 | def zero_module(module): 69 | """ 70 | Zero out the parameters of a module and return it. 71 | """ 72 | for p in module.parameters(): 73 | p.detach().zero_() 74 | return module 75 | 76 | 77 | def scale_module(module, scale): 78 | """ 79 | Scale the parameters of a module and return it. 80 | """ 81 | for p in module.parameters(): 82 | p.detach().mul_(scale) 83 | return module 84 | 85 | 86 | def mean_flat(tensor): 87 | """ 88 | Take the mean over all non-batch dimensions. 89 | """ 90 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 91 | 92 | 93 | def normalization(channels): 94 | """ 95 | Make a standard normalization layer. 96 | 97 | :param channels: number of input channels. 98 | :return: an nn.Module for normalization. 99 | """ 100 | return GroupNorm32(32, channels) 101 | 102 | 103 | def timestep_embedding(timesteps, dim, max_period=10000): 104 | """ 105 | Create sinusoidal timestep embeddings. 106 | 107 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 108 | These may be fractional. 109 | :param dim: the dimension of the output. 110 | :param max_period: controls the minimum frequency of the embeddings. 111 | :return: an [N x dim] Tensor of positional embeddings. 112 | """ 113 | half = dim // 2 114 | freqs = th.exp( 115 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 116 | ).to(device=timesteps.device) 117 | args = timesteps[:, None].float() * freqs[None] 118 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 119 | if dim % 2: 120 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 121 | return embedding 122 | 123 | 124 | def checkpoint(func, inputs, params, flag): 125 | """ 126 | Evaluate a function without caching intermediate activations, allowing for 127 | reduced memory at the expense of extra compute in the backward pass. 128 | 129 | :param func: the function to evaluate. 130 | :param inputs: the argument sequence to pass to `func`. 131 | :param params: a sequence of parameters `func` depends on but does not 132 | explicitly take as arguments. 133 | :param flag: if False, disable gradient checkpointing. 134 | """ 135 | if flag: 136 | args = tuple(inputs) + tuple(params) 137 | return CheckpointFunction.apply(func, len(inputs), *args) 138 | else: 139 | return func(*inputs) 140 | 141 | 142 | class CheckpointFunction(th.autograd.Function): 143 | @staticmethod 144 | def forward(ctx, run_function, length, *args): 145 | ctx.run_function = run_function 146 | ctx.input_tensors = list(args[:length]) 147 | ctx.input_params = list(args[length:]) 148 | with th.no_grad(): 149 | output_tensors = ctx.run_function(*ctx.input_tensors) 150 | return output_tensors 151 | 152 | @staticmethod 153 | def backward(ctx, *output_grads): 154 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 155 | with th.enable_grad(): 156 | # Fixes a bug where the first op in run_function modifies the 157 | # Tensor storage in place, which is not allowed for detach()'d 158 | # Tensors. 159 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 160 | output_tensors = ctx.run_function(*shallow_copies) 161 | input_grads = th.autograd.grad( 162 | output_tensors, 163 | ctx.input_tensors + ctx.input_params, 164 | output_grads, 165 | allow_unused=True, 166 | ) 167 | del ctx.input_tensors 168 | del ctx.input_params 169 | del output_tensors 170 | return (None, None) + input_grads 171 | -------------------------------------------------------------------------------- /SUM/reco_train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch_radon import RadonFanbeam, Radon 3 | import numpy as np 4 | import argparse 5 | from pathlib import Path 6 | from .helper import load_tiff_stack_with_metadata, save_to_tiff_stack 7 | 8 | def para_prepare_parallel(index): 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('--path_proj', type=str, default='/root/autodl-tmp/DDDM/fan_projections.tif', help='Local path of fan beam projection data.') #Reserve a tif file to quickly pass parameters 11 | parser.add_argument('--image_size', type=int, default=736, help='Size of reconstructed image.') # in line with detector number 12 | parser.add_argument('--voxel_size', type=float, default=0.7, help='In-slice voxel size [mm].') 13 | parser.add_argument('--fbp_filter', type=str, default='hann', nargs='?',choices=['ram-lak', 'shepp-logan', 'cosine', 'hamming', 'hann'], help='Filter used for FBP.') 14 | args = parser.parse_args() 15 | 16 | _, metadata = load_tiff_stack_with_metadata(Path(args.path_proj)) 17 | 18 | vox_scaling = 1 / args.voxel_size 19 | bias = 0 # To control the rotation angle of CT image 20 | angles = np.array(metadata['angles'])[:metadata['rotview']+bias] + (np.pi / 2) 21 | 22 | if index == 1: 23 | angles = angles[np.arange(0+bias, 1152+bias, index)] 24 | elif index == 16.5: #.5 means half-scan, integer part means the down-sample factor. 25 | angles = (angles[np.arange(0+bias, 1152+bias, 16)])[np.arange(0,36,1)] 26 | elif index == 8.5: 27 | angles = (angles[np.arange(0+bias, 1152+bias, 8)])[np.arange(0,72,1)] 28 | elif index == 4.5: 29 | angles = (angles[np.arange(0+bias, 1152+bias, 4)])[np.arange(0,144,1)] 30 | elif index == 2.5: 31 | angles = (angles[np.arange(0+bias, 1152+bias, 2)])[np.arange(0,288,1)] 32 | 33 | 34 | # radon = RadonFanbeam(args.image_size, 35 | # angles, 36 | # source_distance=vox_scaling * metadata['dso'], 37 | # det_distance=vox_scaling * metadata['ddo'], 38 | # det_count=736, 39 | # det_spacing=vox_scaling * metadata['du'], 40 | # clip_to_circle=False) 41 | radon = Radon(736, angles=angles, clip_to_circle=True) 42 | return radon 43 | 44 | def run_reco(projections, radon): 45 | # projections = projections[:,range_clip,:] # If FBP results are weired, try uncomment this line. 46 | if(len(projections.shape) == 4): 47 | sino = projections 48 | elif (len(projections.shape) == 3): 49 | sino = torch.flip(projections, dims=[2]) 50 | elif (len(projections.shape) == 2): 51 | sino = torch.flip(projections, dims=[1]) 52 | 53 | filtered_sinogram = radon.filter_sinogram(sino, filter_name='hann') 54 | fbp = 100 * radon.backprojection(filtered_sinogram) 55 | 56 | return fbp 57 | 58 | -------------------------------------------------------------------------------- /SUM/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | import torch.distributed as dist 6 | 7 | 8 | def create_named_schedule_sampler(name, diffusion): 9 | """ 10 | Create a ScheduleSampler from a library of pre-defined samplers. 11 | 12 | :param name: the name of the sampler. 13 | :param diffusion: the diffusion object to sample for. 14 | """ 15 | if name == "uniform": 16 | return UniformSampler(diffusion) 17 | elif name == "loss-second-moment": 18 | return LossSecondMomentResampler(diffusion) 19 | else: 20 | raise NotImplementedError(f"unknown schedule sampler: {name}") 21 | 22 | 23 | class ScheduleSampler(ABC): 24 | """ 25 | A distribution over timesteps in the diffusion process, intended to reduce 26 | variance of the objective. 27 | 28 | By default, samplers perform unbiased importance sampling, in which the 29 | objective's mean is unchanged. 30 | However, subclasses may override sample() to change how the resampled 31 | terms are reweighted, allowing for actual changes in the objective. 32 | """ 33 | 34 | @abstractmethod 35 | def weights(self): 36 | """ 37 | Get a numpy array of weights, one per diffusion step. 38 | 39 | The weights needn't be normalized, but must be positive. 40 | """ 41 | 42 | def sample(self, batch_size, device): 43 | """ 44 | Importance-sample timesteps for a batch. 45 | 46 | :param batch_size: the number of timesteps. 47 | :param device: the torch device to save to. 48 | :return: a tuple (timesteps, weights): 49 | - timesteps: a tensor of timestep indices. 50 | - weights: a tensor of weights to scale the resulting losses. 51 | """ 52 | w = self.weights() 53 | p = w / np.sum(w) 54 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 55 | indices = th.from_numpy(indices_np).long().to(device) 56 | weights_np = 1 / (len(p) * p[indices_np]) 57 | weights = th.from_numpy(weights_np).float().to(device) 58 | return indices, weights 59 | 60 | 61 | class UniformSampler(ScheduleSampler): 62 | def __init__(self, diffusion): 63 | self.diffusion = diffusion 64 | self._weights = np.ones([diffusion.num_timesteps]) 65 | 66 | def weights(self): 67 | return self._weights 68 | 69 | 70 | class LossAwareSampler(ScheduleSampler): 71 | def update_with_local_losses(self, local_ts, local_losses): 72 | """ 73 | Update the reweighting using losses from a model. 74 | 75 | Call this method from each rank with a batch of timesteps and the 76 | corresponding losses for each of those timesteps. 77 | This method will perform synchronization to make sure all of the ranks 78 | maintain the exact same reweighting. 79 | 80 | :param local_ts: an integer Tensor of timesteps. 81 | :param local_losses: a 1D Tensor of losses. 82 | """ 83 | batch_sizes = [ 84 | th.tensor([0], dtype=th.int32, device=local_ts.device) 85 | for _ in range(dist.get_world_size()) 86 | ] 87 | dist.all_gather( 88 | batch_sizes, 89 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 90 | ) 91 | 92 | # Pad all_gather batches to be the maximum batch size. 93 | batch_sizes = [x.item() for x in batch_sizes] 94 | max_bs = max(batch_sizes) 95 | 96 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 97 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 98 | dist.all_gather(timestep_batches, local_ts) 99 | dist.all_gather(loss_batches, local_losses) 100 | timesteps = [ 101 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 102 | ] 103 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 104 | self.update_with_all_losses(timesteps, losses) 105 | 106 | @abstractmethod 107 | def update_with_all_losses(self, ts, losses): 108 | """ 109 | Update the reweighting using losses from a model. 110 | 111 | Sub-classes should override this method to update the reweighting 112 | using losses from the model. 113 | 114 | This method directly updates the reweighting without synchronizing 115 | between workers. It is called by update_with_local_losses from all 116 | ranks with identical arguments. Thus, it should have deterministic 117 | behavior to maintain state across workers. 118 | 119 | :param ts: a list of int timesteps. 120 | :param losses: a list of float losses, one per timestep. 121 | """ 122 | 123 | 124 | class LossSecondMomentResampler(LossAwareSampler): 125 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 126 | self.diffusion = diffusion 127 | self.history_per_term = history_per_term 128 | self.uniform_prob = uniform_prob 129 | self._loss_history = np.zeros( 130 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 131 | ) 132 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 133 | 134 | def weights(self): 135 | if not self._warmed_up(): 136 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 137 | weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) 138 | weights /= np.sum(weights) 139 | weights *= 1 - self.uniform_prob 140 | weights += self.uniform_prob / len(weights) 141 | return weights 142 | 143 | def update_with_all_losses(self, ts, losses): 144 | for t, loss in zip(ts, losses): 145 | if self._loss_counts[t] == self.history_per_term: 146 | # Shift out the oldest loss term. 147 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 148 | self._loss_history[t, -1] = loss 149 | else: 150 | self._loss_history[t, self._loss_counts[t]] = loss 151 | self._loss_counts[t] += 1 152 | 153 | def _warmed_up(self): 154 | return (self._loss_counts == self.history_per_term).all() 155 | -------------------------------------------------------------------------------- /SUM/respace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | 4 | from .gaussian_diffusion import GaussianDiffusion 5 | 6 | 7 | def space_timesteps(num_timesteps, section_counts): 8 | """ 9 | Create a list of timesteps to use from an original diffusion process, 10 | given the number of timesteps we want to take from equally-sized portions 11 | of the original process. 12 | 13 | For example, if there's 300 timesteps and the section counts are [10,15,20] 14 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 15 | are strided to be 15 timesteps, and the final 100 are strided to be 20. 16 | 17 | If the stride is a string starting with "ddim", then the fixed striding 18 | from the DDIM paper is used, and only one section is allowed. 19 | 20 | :param num_timesteps: the number of diffusion steps in the original 21 | process to divide up. 22 | :param section_counts: either a list of numbers, or a string containing 23 | comma-separated numbers, indicating the step count 24 | per section. As a special case, use "ddimN" where N 25 | is a number of steps to use the striding from the 26 | DDIM paper. 27 | :return: a set of diffusion steps from the original process to use. 28 | """ 29 | if isinstance(section_counts, str): 30 | if section_counts.startswith("ddim"): 31 | desired_count = int(section_counts[len("ddim") :]) 32 | for i in range(1, num_timesteps): 33 | if len(range(0, num_timesteps, i)) == desired_count: 34 | return set(range(0, num_timesteps, i)) 35 | raise ValueError( 36 | f"cannot create exactly {num_timesteps} steps with an integer stride" 37 | ) 38 | section_counts = [int(x) for x in section_counts.split(",")] 39 | size_per = num_timesteps // len(section_counts) 40 | extra = num_timesteps % len(section_counts) 41 | start_idx = 0 42 | all_steps = [] 43 | for i, section_count in enumerate(section_counts): 44 | size = size_per + (1 if i < extra else 0) 45 | if size < section_count: 46 | raise ValueError( 47 | f"cannot divide section of {size} steps into {section_count}" 48 | ) 49 | if section_count <= 1: 50 | frac_stride = 1 51 | else: 52 | frac_stride = (size - 1) / (section_count - 1) 53 | cur_idx = 0.0 54 | taken_steps = [] 55 | for _ in range(section_count): 56 | taken_steps.append(start_idx + round(cur_idx)) 57 | cur_idx += frac_stride 58 | all_steps += taken_steps 59 | start_idx += size 60 | return set(all_steps) 61 | 62 | 63 | class SpacedDiffusion(GaussianDiffusion): 64 | """ 65 | A diffusion process which can skip steps in a base diffusion process. 66 | 67 | :param use_timesteps: a collection (sequence or set) of timesteps from the 68 | original diffusion process to retain. 69 | :param kwargs: the kwargs to create the base diffusion process. 70 | """ 71 | 72 | def __init__(self, use_timesteps, **kwargs): 73 | self.use_timesteps = set(use_timesteps) 74 | self.timestep_map = [] 75 | self.original_num_steps = len(kwargs["betas"]) 76 | 77 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa 78 | last_alpha_cumprod = 1.0 79 | new_betas = [] 80 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): 81 | if i in self.use_timesteps: 82 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) 83 | last_alpha_cumprod = alpha_cumprod 84 | self.timestep_map.append(i) 85 | kwargs["betas"] = np.array(new_betas) 86 | super().__init__(**kwargs) 87 | 88 | def p_mean_variance( 89 | self, model, *args, **kwargs 90 | ): # pylint: disable=signature-differs 91 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) 92 | 93 | def training_losses( 94 | self, model, *args, **kwargs 95 | ): # pylint: disable=signature-differs 96 | return super().training_losses(self._wrap_model(model), *args, **kwargs) 97 | 98 | def _wrap_model(self, model): 99 | if isinstance(model, _WrappedModel): 100 | return model 101 | return _WrappedModel( 102 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps 103 | ) 104 | 105 | def _scale_timesteps(self, t): 106 | # Scaling is done by the wrapped model. 107 | return t 108 | 109 | 110 | class _WrappedModel: 111 | def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): 112 | self.model = model 113 | self.timestep_map = timestep_map 114 | self.rescale_timesteps = rescale_timesteps 115 | self.original_num_steps = original_num_steps 116 | 117 | def __call__(self, x, ts, **kwargs): 118 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 119 | new_ts = map_tensor[ts] 120 | if self.rescale_timesteps: 121 | new_ts = new_ts.float() * (1000.0 / self.original_num_steps) 122 | return self.model(x, new_ts, **kwargs) 123 | -------------------------------------------------------------------------------- /SUM/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import inspect 3 | 4 | from . import gaussian_diffusion as gd 5 | from .respace import SpacedDiffusion, space_timesteps 6 | from .unet import SuperResModel, UNetModel 7 | 8 | NUM_CLASSES = 3 9 | 10 | 11 | def model_and_diffusion_defaults(): 12 | """ 13 | Defaults for image training. 14 | """ 15 | return dict( 16 | # image_size=64, 17 | num_channels=128, 18 | num_res_blocks=2, 19 | num_heads=4, 20 | num_heads_upsample=-1, 21 | attention_resolutions="18,9", 22 | dropout=0.3, 23 | learn_sigma=False, 24 | sigma_small=False, 25 | class_cond=False, 26 | diffusion_steps=3, 27 | noise_schedule="cosine", 28 | timestep_respacing="", 29 | use_kl=False, 30 | predict_xstart=True, 31 | rescale_timesteps=False, 32 | rescale_learned_sigmas=False, 33 | use_checkpoint=False, 34 | use_scale_shift_norm=True, 35 | ) 36 | 37 | 38 | def create_model_and_diffusion( 39 | image_size, 40 | class_cond, 41 | learn_sigma, 42 | sigma_small, 43 | num_channels, 44 | num_res_blocks, 45 | num_heads, 46 | num_heads_upsample, 47 | attention_resolutions, 48 | dropout, 49 | diffusion_steps, 50 | noise_schedule, 51 | timestep_respacing, 52 | use_kl, 53 | predict_xstart, 54 | rescale_timesteps, 55 | rescale_learned_sigmas, 56 | use_checkpoint, 57 | use_scale_shift_norm, 58 | ): 59 | model = create_model( 60 | image_size, 61 | num_channels, 62 | num_res_blocks, 63 | learn_sigma=learn_sigma, 64 | class_cond=class_cond, 65 | use_checkpoint=use_checkpoint, 66 | attention_resolutions=attention_resolutions, 67 | num_heads=num_heads, 68 | num_heads_upsample=num_heads_upsample, 69 | use_scale_shift_norm=use_scale_shift_norm, 70 | dropout=dropout, 71 | ) 72 | diffusion = create_gaussian_diffusion( 73 | steps=diffusion_steps, 74 | learn_sigma=learn_sigma, 75 | sigma_small=sigma_small, 76 | noise_schedule=noise_schedule, 77 | use_kl=use_kl, 78 | predict_xstart=predict_xstart, 79 | rescale_timesteps=rescale_timesteps, 80 | rescale_learned_sigmas=rescale_learned_sigmas, 81 | timestep_respacing=timestep_respacing, 82 | ) 83 | return model, diffusion 84 | 85 | 86 | def create_model( 87 | image_size, 88 | num_channels, 89 | num_res_blocks, 90 | learn_sigma, 91 | class_cond, 92 | use_checkpoint, 93 | attention_resolutions, 94 | num_heads, 95 | num_heads_upsample, 96 | use_scale_shift_norm, 97 | dropout, 98 | ): 99 | if image_size == 256 or 288: 100 | channel_mult = (1, 1, 2, 2, 4, 4) 101 | elif image_size == 64: 102 | channel_mult = (1, 2, 3, 4) 103 | elif image_size == 32: 104 | channel_mult = (1, 2, 2, 2) 105 | else: 106 | raise ValueError(f"unsupported image size: {image_size}") 107 | 108 | attention_ds = [] 109 | for res in attention_resolutions.split(","): 110 | attention_ds.append(image_size // int(res)) 111 | 112 | return UNetModel( 113 | in_channels=1, 114 | model_channels=num_channels, 115 | out_channels=(1 if not learn_sigma else 2), 116 | num_res_blocks=num_res_blocks, 117 | attention_resolutions=tuple(attention_ds), 118 | dropout=dropout, 119 | channel_mult=channel_mult, 120 | num_classes=(NUM_CLASSES if class_cond else None), 121 | use_checkpoint=use_checkpoint, 122 | num_heads=num_heads, 123 | num_heads_upsample=num_heads_upsample, 124 | use_scale_shift_norm=use_scale_shift_norm, 125 | ) 126 | 127 | 128 | def sr_model_and_diffusion_defaults(): 129 | res = model_and_diffusion_defaults() 130 | res["large_size"] = 288 131 | res["small_size"] = [36, 736] 132 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 133 | for k in res.copy().keys(): 134 | if k not in arg_names: 135 | del res[k] 136 | return res 137 | 138 | 139 | def sr_create_model_and_diffusion( 140 | large_size, 141 | small_size, 142 | class_cond, 143 | learn_sigma, 144 | num_channels, 145 | num_res_blocks, 146 | num_heads, 147 | num_heads_upsample, 148 | attention_resolutions, 149 | dropout, 150 | diffusion_steps, 151 | noise_schedule, 152 | timestep_respacing, 153 | use_kl, 154 | predict_xstart, 155 | rescale_timesteps, 156 | rescale_learned_sigmas, 157 | use_checkpoint, 158 | use_scale_shift_norm, 159 | ): 160 | model = sr_create_model( 161 | large_size, 162 | small_size, 163 | num_channels, 164 | num_res_blocks, 165 | learn_sigma=learn_sigma, 166 | class_cond=class_cond, 167 | use_checkpoint=use_checkpoint, 168 | attention_resolutions=attention_resolutions, 169 | num_heads=num_heads, 170 | num_heads_upsample=num_heads_upsample, 171 | use_scale_shift_norm=use_scale_shift_norm, 172 | dropout=dropout, 173 | ) 174 | diffusion = create_gaussian_diffusion( 175 | steps=diffusion_steps, 176 | learn_sigma=learn_sigma, 177 | noise_schedule=noise_schedule, 178 | use_kl=use_kl, 179 | predict_xstart=predict_xstart, 180 | rescale_timesteps=rescale_timesteps, 181 | rescale_learned_sigmas=rescale_learned_sigmas, 182 | timestep_respacing=timestep_respacing, 183 | ) 184 | return model, diffusion 185 | 186 | 187 | def sr_create_model( 188 | large_size, 189 | small_size, 190 | num_channels, 191 | num_res_blocks, 192 | learn_sigma, 193 | class_cond, 194 | use_checkpoint, 195 | attention_resolutions, 196 | num_heads, 197 | num_heads_upsample, 198 | use_scale_shift_norm, 199 | dropout, 200 | ): 201 | _ = small_size # hack to prevent unused variable 202 | 203 | if large_size == 256 or 288: 204 | channel_mult = (1, 1, 2, 2, 4, 4) 205 | elif large_size == 64: 206 | channel_mult = (1, 2, 3, 4) 207 | else: 208 | raise ValueError(f"unsupported large size: {large_size}") 209 | 210 | attention_ds = [] 211 | for res in attention_resolutions.split(","): 212 | attention_ds.append(large_size // int(res)) 213 | 214 | return SuperResModel( 215 | in_channels=1, 216 | model_channels=num_channels, 217 | out_channels=(1 if not learn_sigma else 2), 218 | num_res_blocks=num_res_blocks, 219 | attention_resolutions=tuple(attention_ds), 220 | dropout=dropout, 221 | channel_mult=channel_mult, 222 | num_classes=(NUM_CLASSES if class_cond else None), 223 | use_checkpoint=use_checkpoint, 224 | num_heads=num_heads, 225 | num_heads_upsample=num_heads_upsample, 226 | use_scale_shift_norm=use_scale_shift_norm, 227 | ) 228 | 229 | 230 | def create_gaussian_diffusion( 231 | *, 232 | steps=4000, 233 | learn_sigma=False, 234 | sigma_small=False, 235 | noise_schedule="cosine", 236 | use_kl=False, 237 | predict_xstart=True, 238 | rescale_timesteps=False, 239 | rescale_learned_sigmas=False, 240 | timestep_respacing="", 241 | ): 242 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 243 | if use_kl: 244 | loss_type = gd.LossType.RESCALED_KL 245 | elif rescale_learned_sigmas: 246 | loss_type = gd.LossType.RESCALED_MSE 247 | else: 248 | loss_type = gd.LossType.MSE 249 | if not timestep_respacing: 250 | timestep_respacing = [steps] 251 | return SpacedDiffusion( 252 | use_timesteps=space_timesteps(steps, timestep_respacing), 253 | betas=betas, 254 | model_mean_type=( 255 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 256 | ), 257 | model_var_type=( 258 | ( 259 | gd.ModelVarType.FIXED_LARGE 260 | if not sigma_small 261 | else gd.ModelVarType.FIXED_SMALL 262 | ) 263 | if not learn_sigma 264 | else gd.ModelVarType.LEARNED_RANGE 265 | ), 266 | loss_type=loss_type, 267 | rescale_timesteps=rescale_timesteps, 268 | ) 269 | 270 | 271 | def add_dict_to_argparser(parser, default_dict): 272 | for k, v in default_dict.items(): 273 | v_type = type(v) 274 | if v is None: 275 | v_type = str 276 | elif isinstance(v, bool): 277 | v_type = str2bool 278 | parser.add_argument(f"--{k}", default=v, type=v_type) 279 | 280 | 281 | def args_to_dict(args, keys): 282 | return {k: getattr(args, k) for k in keys} 283 | 284 | 285 | def str2bool(v): 286 | """ 287 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 288 | """ 289 | if isinstance(v, bool): 290 | return v 291 | if v.lower() in ("yes", "true", "t", "y", "1"): 292 | return True 293 | elif v.lower() in ("no", "false", "f", "n", "0"): 294 | return False 295 | else: 296 | raise argparse.ArgumentTypeError("boolean value expected") 297 | -------------------------------------------------------------------------------- /SUM/scripts/image_nll.py: -------------------------------------------------------------------------------- 1 | """ 2 | Approximate the bits/dimension for an image model. 3 | """ 4 | 5 | import argparse 6 | import os 7 | 8 | import numpy as np 9 | import torch.distributed as dist 10 | 11 | from improved_diffusion import dist_util, logger 12 | from improved_diffusion.image_datasets import load_data 13 | from improved_diffusion.script_util import ( 14 | model_and_diffusion_defaults, 15 | create_model_and_diffusion, 16 | add_dict_to_argparser, 17 | args_to_dict, 18 | ) 19 | 20 | 21 | def main(): 22 | args = create_argparser().parse_args() 23 | 24 | dist_util.setup_dist() 25 | logger.configure() 26 | 27 | logger.log("creating model and diffusion...") 28 | model, diffusion = create_model_and_diffusion( 29 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 30 | ) 31 | model.load_state_dict( 32 | dist_util.load_state_dict(args.model_path, map_location="cpu") 33 | ) 34 | model.to(dist_util.dev()) 35 | model.eval() 36 | 37 | logger.log("creating data loader...") 38 | data = load_data( 39 | data_dir=args.data_dir, 40 | batch_size=args.batch_size, 41 | image_size=args.image_size, 42 | class_cond=args.class_cond, 43 | deterministic=True, 44 | ) 45 | 46 | logger.log("evaluating...") 47 | run_bpd_evaluation(model, diffusion, data, args.num_samples, args.clip_denoised) 48 | 49 | 50 | def run_bpd_evaluation(model, diffusion, data, num_samples, clip_denoised): 51 | all_bpd = [] 52 | all_metrics = {"vb": [], "mse": [], "xstart_mse": []} 53 | num_complete = 0 54 | while num_complete < num_samples: 55 | batch, model_kwargs = next(data) 56 | batch = batch.to(dist_util.dev()) 57 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 58 | minibatch_metrics = diffusion.calc_bpd_loop( 59 | model, batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs 60 | ) 61 | 62 | for key, term_list in all_metrics.items(): 63 | terms = minibatch_metrics[key].mean(dim=0) / dist.get_world_size() 64 | dist.all_reduce(terms) 65 | term_list.append(terms.detach().cpu().numpy()) 66 | 67 | total_bpd = minibatch_metrics["total_bpd"] 68 | total_bpd = total_bpd.mean() / dist.get_world_size() 69 | dist.all_reduce(total_bpd) 70 | all_bpd.append(total_bpd.item()) 71 | num_complete += dist.get_world_size() * batch.shape[0] 72 | 73 | logger.log(f"done {num_complete} samples: bpd={np.mean(all_bpd)}") 74 | 75 | if dist.get_rank() == 0: 76 | for name, terms in all_metrics.items(): 77 | out_path = os.path.join(logger.get_dir(), f"{name}_terms.npz") 78 | logger.log(f"saving {name} terms to {out_path}") 79 | np.savez(out_path, np.mean(np.stack(terms), axis=0)) 80 | 81 | dist.barrier() 82 | logger.log("evaluation complete") 83 | 84 | 85 | def create_argparser(): 86 | defaults = dict( 87 | data_dir="", clip_denoised=True, num_samples=1000, batch_size=1, model_path="" 88 | ) 89 | defaults.update(model_and_diffusion_defaults()) 90 | parser = argparse.ArgumentParser() 91 | add_dict_to_argparser(parser, defaults) 92 | return parser 93 | 94 | 95 | if __name__ == "__main__": 96 | main() 97 | -------------------------------------------------------------------------------- /SUM/scripts/multi_super_res.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of samples from a super resolution model, given a batch 3 | of samples from a regular model from image_sample.py. 4 | """ 5 | import matplotlib.pyplot as plt 6 | import argparse 7 | import os 8 | import torch.nn.functional as F 9 | from improved_diffusion.image_datasets import trans_data_preprocess 10 | import blobfile as bf 11 | import numpy as np 12 | import torch as th 13 | import torch.distributed as dist 14 | from skimage.metrics import structural_similarity as ssim 15 | from skimage.metrics import peak_signal_noise_ratio as psnr 16 | import matplotlib.pyplot as plt 17 | import lpips 18 | from improved_diffusion.reco_train import run_reco, para_prepare_parallel 19 | 20 | 21 | from improved_diffusion import dist_util, logger 22 | from improved_diffusion.script_util import ( 23 | sr_model_and_diffusion_defaults, 24 | sr_create_model_and_diffusion, 25 | args_to_dict, 26 | add_dict_to_argparser, 27 | ) 28 | 29 | def p_sample_loop_super_res(model, batch_size, model_kwargs): 30 | 31 | with th.no_grad(): 32 | x_in = F.interpolate(model_kwargs["low_res"], [288,736], mode="nearest") 33 | x2 = model(x_in, timesteps=th.tensor([2]).to("cuda"), **model_kwargs) 34 | x1 = model(x2, timesteps=th.tensor([1]).to("cuda"), **model_kwargs) 35 | x0 = model(x1, timesteps=th.tensor([0]).to("cuda"), **model_kwargs) 36 | x_out = x0 37 | return x_out, x2 38 | 39 | 40 | def main(): 41 | num = 800 # Total generated picture num 42 | args = create_argparser().parse_args() 43 | 44 | dist_util.setup_dist() 45 | logger.configure() 46 | 47 | logger.log("creating model...") 48 | model, diffusion = sr_create_model_and_diffusion( 49 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 50 | ) 51 | model.load_state_dict( 52 | dist_util.load_state_dict(args.model_path, map_location="cuda") 53 | ) 54 | model.to(dist_util.dev()) 55 | model.eval() 56 | 57 | logger.log("loading data...") 58 | 59 | data = load_superres_data(args.batch_size) 60 | 61 | logger.log("creating samples...") 62 | 63 | MSE, SSIM, PSNR, LPIP = [], [], [], [] 64 | lpip_loss = lpips.LPIPS(net="alex").to(dist_util.dev()) 65 | 66 | l2loss = th.nn.MSELoss().to(dist_util.dev()) 67 | 68 | 69 | radon_288_736 = para_prepare_parallel(2.5) 70 | radon_72_736 = para_prepare_parallel(8.5) 71 | radon_36_736 = para_prepare_parallel(16.5) 72 | helper = {"fbp_para_288_736": radon_288_736, "fbp_para_36_736": radon_36_736, "fbp_para_72_736": radon_72_736} 73 | 74 | for i in range(0, num//args.batch_size):# 75 | model_kwargs = next(data) 76 | raw_img = model_kwargs.pop('raw_img').to("cuda") 77 | index = model_kwargs.pop('index') 78 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 79 | model_kwargs["fbp_para_36_736"] = radon_36_736 80 | model_kwargs["fbp_para_288_736"] = radon_288_736 81 | 82 | sample_fn = p_sample_loop_super_res 83 | sample, sample_72_288 = sample_fn( 84 | model, 85 | (args.batch_size, 1, 288, 736), #args.large_size, args.large_size 86 | # clip_denoised=args.clip_denoised, 87 | model_kwargs=model_kwargs, 88 | ) 89 | 90 | model_72_sino = F.interpolate(sample_72_288, [72, 736], mode="nearest") 91 | model_72_fbp = run_reco(model_72_sino + 1., helper["fbp_para_72_736"])[:,:,112:624,112:624] 92 | model_72_fbp_npy = model_72_fbp.cpu().detach().numpy() 93 | 94 | model_output_fbp = run_reco(sample + 1., helper["fbp_para_288_736"])[:,:,112:624,112:624] 95 | target_fbp = run_reco(raw_img + 1., helper["fbp_para_288_736"])[:,:,112:624,112:624] 96 | output_fbp_npy = model_output_fbp.cpu().detach().numpy() 97 | 98 | for j in range(0, args.batch_size): 99 | 100 | l2loss_value = l2loss(model_output_fbp[j], target_fbp[j]).item() 101 | print("index:", index[j], "MSELoss:", l2loss_value) 102 | MSE.append(l2loss_value) 103 | 104 | raw_npy = target_fbp.cpu().detach().numpy() 105 | ssim_value = ssim(np.squeeze(output_fbp_npy[j]),np.squeeze( raw_npy[j]), data_range = raw_npy[j].max() - raw_npy[j].min()) 106 | psnr_value = psnr(np.squeeze(output_fbp_npy[j]),np.squeeze( raw_npy[j]), data_range = raw_npy[j].max() - raw_npy[j].min()) 107 | print("index:", index[j], "SSIM:", ssim_value) 108 | SSIM.append(ssim_value) 109 | PSNR.append(psnr_value) 110 | 111 | lpip_value = lpip_loss(model_output_fbp[j], target_fbp[j]) 112 | print("lpips:", lpip_value.item()) 113 | LPIP.append(lpip_value.item()) 114 | 115 | print("mse mean:", np.mean(MSE)) 116 | print("ssim mean:", np.mean(SSIM)) 117 | print("psnr mean:", np.mean(PSNR)) 118 | print("lpip mean:", np.mean(LPIP)) 119 | 120 | 121 | def create_argparser(): 122 | defaults = dict( 123 | clip_denoised=False, 124 | num_samples=1, 125 | batch_size=2, 126 | use_ddim=True, 127 | base_samples="", 128 | model_path='/root/autodl-tmp/TASK8/improved_diffusion/model_save_FULL/ema_0.9999_036000.pt', 129 | ) 130 | print(defaults["model_path"]) 131 | defaults.update(sr_model_and_diffusion_defaults()) 132 | parser = argparse.ArgumentParser() 133 | add_dict_to_argparser(parser, defaults) 134 | return parser 135 | 136 | 137 | def load_superres_data(batch_size): 138 | 139 | data = trans_data_preprocess(batch_size = batch_size, shuffle = False, num_workers = 8) 140 | 141 | for large_batch, model_kwargs in data: 142 | 143 | large_batch = large_batch[:,:,np.arange(0, 576, 2),:] 144 | 145 | model_kwargs["raw_img"] = large_batch 146 | model_kwargs["low_res"] = F.interpolate(large_batch, [36, 736], mode="nearest") #36->72 147 | 148 | res = dict(low_res=model_kwargs["low_res"], raw_img=model_kwargs["raw_img"],index=model_kwargs["index"] ) 149 | yield res 150 | 151 | 152 | if __name__ == "__main__": 153 | main() -------------------------------------------------------------------------------- /SUM/scripts/super_res_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a single batch of samples from a super resolution model. 3 | """ 4 | import matplotlib.pyplot as plt 5 | import argparse 6 | import os 7 | import torch.nn.functional as F 8 | from improved_diffusion.image_datasets import improved_data_preprocess_val 9 | import blobfile as bf 10 | import numpy as np 11 | import torch as th 12 | import torch.distributed as dist 13 | from skimage.metrics import structural_similarity as ssim 14 | import lpips 15 | from improved_diffusion.reco_train import run_reco, para_prepare_parallel, para_prepare 16 | 17 | from improved_diffusion import dist_util, logger 18 | from improved_diffusion.script_util import ( 19 | sr_model_and_diffusion_defaults, 20 | sr_create_model_and_diffusion, 21 | args_to_dict, 22 | add_dict_to_argparser, 23 | ) 24 | 25 | 26 | model_name = "ema_0.9999_028500" 27 | parallel = True 28 | 29 | 30 | def down_up(input,down_size,mode): 31 | output = F.interpolate(input, down_size, mode="nearest") 32 | output = F.interpolate(output, [288,736], mode=mode) 33 | return output 34 | 35 | 36 | def p_sample_loop_super_res(model, batch_size, model_kwargs, raw_img): 37 | 38 | with th.no_grad(): 39 | x_in = F.interpolate(model_kwargs["low_res"], [288,736], mode="nearest") 40 | x2 = model(x_in, timesteps=th.tensor([2]).to("cuda"), **model_kwargs) 41 | x1 = model(x2, timesteps=th.tensor([1]).to("cuda"), **model_kwargs) 42 | x0 = model(x1, timesteps=th.tensor([0]).to("cuda"), **model_kwargs) 43 | x_out = x0 44 | return x_out 45 | 46 | 47 | def main(): 48 | args = create_argparser().parse_args() 49 | 50 | dist_util.setup_dist() 51 | logger.configure() 52 | 53 | logger.log("creating model...") 54 | model, diffusion = sr_create_model_and_diffusion( 55 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 56 | ) 57 | model.load_state_dict( 58 | dist_util.load_state_dict(args.model_path, map_location="cpu") 59 | ) 60 | model.to(dist_util.dev()) 61 | model.eval() 62 | 63 | logger.log("loading data...") 64 | data = load_superres_data(args.batch_size, small_size=[36,736], class_cond=False) 65 | 66 | logger.log("creating samples...") 67 | all_images = [] 68 | 69 | if parallel: 70 | radon_288_736 = para_prepare_parallel(2.5) 71 | radon_36_736 = para_prepare_parallel(16.5) 72 | 73 | helper = {"fbp_para_288_736": radon_288_736, "fbp_para_36_736": radon_36_736} 74 | 75 | while len(all_images) * args.batch_size < args.num_samples: 76 | model_kwargs = next(data) 77 | raw_img = model_kwargs.pop('raw_img').to("cuda") 78 | raw_img = down_up(raw_img,[288,736],"nearest") 79 | 80 | 81 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 82 | model_kwargs["fbp_para_36_736"] = radon_36_736 83 | model_kwargs["fbp_para_288_736"] = radon_288_736 84 | 85 | input_fbp = run_reco(th.flip(model_kwargs['low_res'].to("cuda") + 1., dims=[3]), helper["fbp_para_36_736"])[:,:,112:624,112:624] 86 | input_npy = input_fbp.squeeze().cpu().detach().numpy() 87 | plt.imshow(input_npy, cmap=plt.cm.gray) 88 | 89 | 90 | sample_fn = p_sample_loop_super_res 91 | sample = sample_fn( 92 | model, 93 | (args.batch_size, 1, 288, 736), #args.large_size, args.large_size 94 | # clip_denoised=args.clip_denoised, 95 | model_kwargs=model_kwargs, 96 | raw_img = raw_img, 97 | ) 98 | 99 | 100 | model_output_fbp = run_reco(th.flip(sample + 1.,dims=[3]), helper["fbp_para_288_736"])[:,:,112:624,112:624] 101 | target_fbp = run_reco(th.flip(raw_img + 1., dims=[3]), helper["fbp_para_288_736"])[:,:,112:624,112:624] 102 | 103 | target_npy = target_fbp.squeeze().cpu().detach().numpy() 104 | plt.imshow(target_npy, cmap=plt.cm.gray) 105 | 106 | npy = np.squeeze(model_output_fbp.cpu().detach().numpy()) 107 | print("mean: ", np.mean(npy)) 108 | print("std: ", np.std(npy)) 109 | print("max: ", np.max(npy)) 110 | print("min: ", np.min(npy)) 111 | 112 | raw_npy = target_fbp.squeeze().cpu().detach().numpy() 113 | print("SSIM:", ssim(npy, raw_npy,data_range=raw_npy.max()-raw_npy.min())) 114 | 115 | l2loss = th.nn.MSELoss().to(dist_util.dev()) 116 | print("MSELoss:", l2loss(model_output_fbp, target_fbp.to(dist_util.dev())).item()) 117 | 118 | lpip_loss = lpips.LPIPS(net="alex").to(dist_util.dev()) 119 | lpip_value = lpip_loss(model_output_fbp, target_fbp.to(dist_util.dev())) 120 | print("lpips:", lpip_value.item()) 121 | 122 | break 123 | 124 | 125 | def create_argparser(): 126 | defaults = dict( 127 | clip_denoised=False, 128 | num_samples=1, 129 | batch_size=1, 130 | use_ddim=True, 131 | base_samples="", 132 | model_path=f'/root/autodl-tmp/improved_diffusion/model_save_FULL/{model_name}.pt', # model065000 ema_0.9999_060000 133 | ) 134 | print(defaults["model_path"]) 135 | defaults.update(sr_model_and_diffusion_defaults()) 136 | parser = argparse.ArgumentParser() 137 | add_dict_to_argparser(parser, defaults) 138 | return parser 139 | 140 | def load_superres_data(batch_size, small_size, class_cond=False): 141 | 142 | data = improved_data_preprocess_val(batch_size = batch_size, shuffle = True, num_workers = 8) 143 | radon_576_736 = para_prepare_parallel(2) 144 | 145 | for large_batch, model_kwargs in data: 146 | if parallel: 147 | _576_fbp = run_reco(th.flip(large_batch[:,:,np.arange(0, 1152, 2),:] + 1., dims=[3]).to("cuda"), radon_576_736)[:,:,112:624,112:624] 148 | 149 | 150 | plt.imshow(_576_fbp.squeeze().cpu().detach().numpy(), cmap=plt.cm.gray) 151 | 152 | large_batch = large_batch[:,:,np.arange(0, 576, 2),:] 153 | npy = large_batch.squeeze().cpu().detach().numpy() 154 | print("mean: ", np.mean(npy)) 155 | print("std: ", np.std(npy)) 156 | print("max: ", np.max(npy)) 157 | print("min: ", np.min(npy)) 158 | 159 | model_kwargs["raw_img"] = large_batch 160 | model_kwargs["low_res"] = F.interpolate(large_batch, [36, 736], mode="nearest") 161 | 162 | res = dict(low_res=model_kwargs["low_res"], raw_img=model_kwargs["raw_img"]) 163 | yield res 164 | 165 | 166 | if __name__ == "__main__": 167 | main() 168 | -------------------------------------------------------------------------------- /SUM/scripts/super_res_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a super-resolution model. 3 | """ 4 | 5 | import argparse 6 | import numpy as np 7 | import torch.nn.functional as F 8 | from improved_diffusion import dist_util, logger 9 | from improved_diffusion.image_datasets import improved_data_preprocess 10 | from improved_diffusion.resample import create_named_schedule_sampler 11 | from improved_diffusion.script_util import ( 12 | sr_model_and_diffusion_defaults, 13 | sr_create_model_and_diffusion, 14 | args_to_dict, 15 | add_dict_to_argparser, 16 | ) 17 | from improved_diffusion.train_util import TrainLoop 18 | 19 | 20 | def main(): 21 | args = create_argparser().parse_args() 22 | 23 | dist_util.setup_dist() 24 | logger.configure() 25 | 26 | logger.log("creating model...") 27 | model, diffusion = sr_create_model_and_diffusion( 28 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 29 | ) 30 | model.to(dist_util.dev()) 31 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 32 | 33 | logger.log("creating data loader...") 34 | 35 | data = load_superres_data( 36 | args.batch_size, 37 | large_size=args.large_size, 38 | small_size=args.small_size, 39 | class_cond=args.class_cond, 40 | ) 41 | 42 | logger.log("training...") 43 | print(args) 44 | 45 | # num_param = sum(p.numel() for p in model.parameters() if p.requires_grad) 46 | # print('Number of parameters in model: {}'.format(num_param)) 47 | # exit() 48 | TrainLoop( 49 | model=model, 50 | diffusion=diffusion, 51 | data=data, 52 | batch_size=args.batch_size, 53 | microbatch=args.microbatch, 54 | lr=args.lr, 55 | ema_rate=args.ema_rate, 56 | log_interval=args.log_interval, 57 | save_interval=args.save_interval, 58 | resume_checkpoint=args.resume_checkpoint, 59 | use_fp16=args.use_fp16, 60 | fp16_scale_growth=args.fp16_scale_growth, 61 | schedule_sampler=schedule_sampler, 62 | weight_decay=args.weight_decay, 63 | lr_anneal_steps=args.lr_anneal_steps, 64 | ).run_loop() 65 | 66 | 67 | def load_superres_data(batch_size, large_size, small_size, class_cond=False): 68 | 69 | print("class_cond:", class_cond) 70 | data = improved_data_preprocess(batch_size = batch_size, shuffle = True, num_workers = 8) 71 | for large_batch, model_kwargs in data: 72 | 73 | large_batch = large_batch[:,:,np.arange(0, 576, 2),:] # 288 half 74 | model_kwargs["low_res"] = F.interpolate(large_batch, size=[36,736], mode='nearest') # 36 half 75 | yield large_batch, model_kwargs 76 | 77 | 78 | def create_argparser(): 79 | defaults = dict( 80 | data_dir="", 81 | schedule_sampler="loss-second-moment", 82 | lr=1e-4, 83 | weight_decay=0.0, 84 | lr_anneal_steps=0, 85 | batch_size=2, 86 | microbatch=-1, 87 | ema_rate="0.9999", 88 | log_interval=150, 89 | save_interval=1500, 90 | resume_checkpoint="",# /root/autodl-tmp/improved_diffusion/model_save_RES/model065000.pt 91 | use_fp16=False, 92 | fp16_scale_growth=1e-3, 93 | ) 94 | defaults.update(sr_model_and_diffusion_defaults()) 95 | parser = argparse.ArgumentParser() 96 | add_dict_to_argparser(parser, defaults) 97 | return parser 98 | 99 | 100 | if __name__ == "__main__": 101 | main() 102 | -------------------------------------------------------------------------------- /SUM/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import numpy as np 7 | import torch as th 8 | import torch.distributed as dist 9 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 10 | from torch.optim import AdamW 11 | from pytorch_msssim import ms_ssim, MS_SSIM, SSIM 12 | from .reco_train import para_prepare_parallel, para_prepare 13 | 14 | from . import dist_util, logger 15 | from .fp16_util import ( 16 | make_master_params, 17 | master_params_to_model_params, 18 | model_grads_to_master_grads, 19 | unflatten_master_params, 20 | zero_grad, 21 | ) 22 | from .nn import update_ema 23 | from .resample import LossAwareSampler, UniformSampler 24 | import torch.nn.functional as F 25 | 26 | import lpips 27 | 28 | # For ImageNet experiments, this was a good default value. 29 | # We found that the lg_loss_scale quickly climbed to 30 | # 20-21 within the first ~1K steps of training. 31 | INITIAL_LOG_LOSS_SCALE = 20.0 32 | parallel = True 33 | 34 | 35 | class TrainLoop: 36 | def __init__( 37 | self, 38 | *, 39 | model, 40 | diffusion, 41 | data, 42 | batch_size, 43 | microbatch, 44 | lr, 45 | ema_rate, 46 | log_interval, 47 | save_interval, 48 | resume_checkpoint, 49 | use_fp16=False, 50 | fp16_scale_growth=1e-3, 51 | schedule_sampler=None, 52 | weight_decay=0.0, 53 | lr_anneal_steps=0, 54 | ): 55 | self.model = model 56 | self.diffusion = diffusion 57 | self.data = data 58 | self.batch_size = batch_size 59 | self.microbatch = microbatch if microbatch > 0 else batch_size 60 | self.lr = lr 61 | self.ema_rate = ( 62 | [ema_rate] 63 | if isinstance(ema_rate, float) 64 | else [float(x) for x in ema_rate.split(",")] 65 | ) 66 | self.log_interval = log_interval 67 | self.save_interval = save_interval 68 | self.resume_checkpoint = resume_checkpoint 69 | self.use_fp16 = use_fp16 70 | self.fp16_scale_growth = fp16_scale_growth 71 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 72 | self.weight_decay = weight_decay 73 | self.lr_anneal_steps = lr_anneal_steps 74 | 75 | self.step = 0 76 | self.resume_step = 0 77 | self.global_batch = self.batch_size * dist.get_world_size() 78 | 79 | self.model_params = list(self.model.parameters()) 80 | self.master_params = self.model_params 81 | self.lg_loss_scale = INITIAL_LOG_LOSS_SCALE 82 | self.sync_cuda = th.cuda.is_available() 83 | 84 | self._load_and_sync_parameters() 85 | if self.use_fp16: 86 | self._setup_fp16() 87 | 88 | self.opt = AdamW(self.master_params, lr=self.lr, weight_decay=self.weight_decay) 89 | if self.resume_step: 90 | self._load_optimizer_state() 91 | # Model was resumed, either due to a restart or a checkpoint 92 | # being specified at the command line. 93 | self.ema_params = [ 94 | self._load_ema_parameters(rate) for rate in self.ema_rate 95 | ] 96 | else: 97 | self.ema_params = [ 98 | copy.deepcopy(self.master_params) for _ in range(len(self.ema_rate)) 99 | ] 100 | 101 | if th.cuda.is_available(): 102 | self.use_ddp = True 103 | self.ddp_model = DDP( 104 | self.model, 105 | device_ids=[dist_util.dev()], 106 | output_device=dist_util.dev(), 107 | broadcast_buffers=False, 108 | bucket_cap_mb=128, 109 | find_unused_parameters=False, 110 | ) 111 | else: 112 | if dist.get_world_size() > 1: 113 | logger.warn( 114 | "Distributed training requires CUDA. " 115 | "Gradients will not be synchronized properly!" 116 | ) 117 | self.use_ddp = False 118 | self.ddp_model = self.model 119 | 120 | def _load_and_sync_parameters(self): 121 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 122 | 123 | if resume_checkpoint: 124 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 125 | if dist.get_rank() == 0: 126 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 127 | self.model.load_state_dict( 128 | dist_util.load_state_dict( 129 | resume_checkpoint, map_location=dist_util.dev() 130 | ) 131 | ) 132 | 133 | dist_util.sync_params(self.model.parameters()) 134 | 135 | def _load_ema_parameters(self, rate): 136 | ema_params = copy.deepcopy(self.master_params) 137 | 138 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 139 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 140 | if ema_checkpoint: 141 | if dist.get_rank() == 0: 142 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 143 | state_dict = dist_util.load_state_dict( 144 | ema_checkpoint, map_location=dist_util.dev() 145 | ) 146 | ema_params = self._state_dict_to_master_params(state_dict) 147 | 148 | dist_util.sync_params(ema_params) 149 | return ema_params 150 | 151 | def _load_optimizer_state(self): 152 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 153 | opt_checkpoint = bf.join( 154 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 155 | ) 156 | if bf.exists(opt_checkpoint): 157 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 158 | state_dict = dist_util.load_state_dict( 159 | opt_checkpoint, map_location=dist_util.dev() 160 | ) 161 | self.opt.load_state_dict(state_dict) 162 | 163 | def _setup_fp16(self): 164 | self.master_params = make_master_params(self.model_params) 165 | self.model.convert_to_fp16() 166 | 167 | def run_loop(self): 168 | lpip_loss = lpips.LPIPS(net="alex").to(dist_util.dev()) ############## 169 | ssim_loss = SSIM(win_size=7, win_sigma=1.5, data_range=1, size_average=False, channel=1) 170 | if parallel: 171 | radon_288_736 = para_prepare_parallel(2.5) 172 | radon_144_736 = para_prepare_parallel(4.5) 173 | radon_72_736 = para_prepare_parallel(8.5) 174 | radon_36_736 = para_prepare_parallel(16.5) 175 | else: 176 | radon_288_736 = para_prepare(2.5) 177 | radon_36_736 = para_prepare(16.5) 178 | helper = {"fbp_para_288_736": radon_288_736, "fbp_para_36_736": radon_36_736, "fbp_para_72_736": radon_72_736, "fbp_para_144_736": radon_144_736} ######################### "fbp_para_36_512": radon_36_51 179 | 180 | while ( 181 | not self.lr_anneal_steps 182 | or self.step + self.resume_step < self.lr_anneal_steps 183 | ): 184 | batch, cond = next(self.data) 185 | 186 | timestep = np.random.randint(low=3, high=None, size=None, dtype='l') 187 | t = th.tensor([timestep,timestep]).to("cuda") 188 | if timestep == 2: 189 | cond["x_t"] = F.interpolate(F.interpolate(batch, (36, 736), mode="nearest"), (288, 736), mode="nearest") 190 | elif timestep == 1: 191 | cond["x_t"] = F.interpolate(F.interpolate(batch, (72, 736), mode="nearest"), (288, 736), mode="nearest") 192 | elif timestep == 0: 193 | cond["x_t"] = F.interpolate(F.interpolate(batch, (144, 736), mode="nearest"), (288, 736), mode="nearest") 194 | model_output = self.run_step(batch, cond, t, ssim_loss, lpip_loss, helper) 195 | 196 | # for i in range(2,-1,-1): 197 | # t = th.tensor([i,i]).to("cuda") 198 | # if i == 2: 199 | # cond["x_t"] = cond["low_res"] 200 | # model_output = self.run_step(batch, cond, t, ssim_loss, lpip_loss, helper) 201 | # else: 202 | # cond["x_t"] = model_output 203 | # model_output = self.run_step(batch, cond, t, ssim_loss, lpip_loss, helper) 204 | 205 | 206 | 207 | 208 | if self.step % self.log_interval == 0: 209 | logger.dumpkvs() 210 | if self.step % self.save_interval == 0: 211 | self.save() 212 | # Run for a finite amount of time in integration tests. 213 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 214 | return 215 | self.step += 1 216 | # Save the last checkpoint if it wasn't already saved. 217 | if (self.step - 1) % self.save_interval != 0: 218 | self.save() 219 | 220 | def run_step(self, batch, cond, t, ssim_loss=None, lpip_loss=None, helper=None): 221 | model_output = self.forward_backward(batch, cond, t, ssim_loss, lpip_loss, helper) 222 | if self.use_fp16: 223 | self.optimize_fp16() 224 | else: 225 | self.optimize_normal() 226 | self.log_step() 227 | return model_output 228 | 229 | def forward_backward(self, batch, cond, t, ssim_loss=None, lpip_loss=None, helper=None): 230 | zero_grad(self.model_params) 231 | for i in range(0, batch.shape[0], self.microbatch): 232 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 233 | micro_cond = { 234 | k: v[i : i + self.microbatch].to(dist_util.dev()) 235 | for k, v in cond.items() 236 | } 237 | last_batch = (i + self.microbatch) >= batch.shape[0] 238 | _, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 239 | 240 | compute_losses = functools.partial( 241 | self.diffusion.training_losses, 242 | self.ddp_model, 243 | micro, 244 | t, 245 | model_kwargs=micro_cond, 246 | ssim_loss = ssim_loss, 247 | lpip_loss = lpip_loss, 248 | helper = helper 249 | ) 250 | 251 | if last_batch or not self.use_ddp: 252 | losses, model_output = compute_losses() 253 | else: 254 | with self.ddp_model.no_sync(): 255 | losses, model_output = compute_losses() 256 | 257 | if isinstance(self.schedule_sampler, LossAwareSampler): 258 | self.schedule_sampler.update_with_local_losses( 259 | t, losses["loss"].detach() 260 | ) 261 | 262 | loss = (losses["loss"] * weights).mean() 263 | log_loss_dict( 264 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 265 | ) 266 | if self.use_fp16: 267 | loss_scale = 2 ** self.lg_loss_scale 268 | (loss * loss_scale).backward() 269 | else: 270 | loss.backward() 271 | 272 | model_output_copy = model_output.data 273 | return model_output_copy 274 | 275 | def optimize_fp16(self): 276 | if any(not th.isfinite(p.grad).all() for p in self.model_params): 277 | self.lg_loss_scale -= 1 278 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 279 | return 280 | 281 | model_grads_to_master_grads(self.model_params, self.master_params) 282 | self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale)) 283 | self._log_grad_norm() 284 | self._anneal_lr() 285 | self.opt.step() 286 | for rate, params in zip(self.ema_rate, self.ema_params): 287 | update_ema(params, self.master_params, rate=rate) 288 | master_params_to_model_params(self.model_params, self.master_params) 289 | self.lg_loss_scale += self.fp16_scale_growth 290 | 291 | def optimize_normal(self): 292 | self._log_grad_norm() 293 | self._anneal_lr() 294 | self.opt.step() 295 | for rate, params in zip(self.ema_rate, self.ema_params): 296 | update_ema(params, self.master_params, rate=rate) 297 | 298 | def _log_grad_norm(self): 299 | sqsum = 0.0 300 | for p in self.master_params: 301 | sqsum += (p.grad ** 2).sum().item() 302 | logger.logkv_mean("grad_norm", np.sqrt(sqsum)) 303 | 304 | def _anneal_lr(self): 305 | if not self.lr_anneal_steps: 306 | return 307 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 308 | lr = self.lr * (1 - frac_done) 309 | for param_group in self.opt.param_groups: 310 | param_group["lr"] = lr 311 | 312 | def log_step(self): 313 | logger.logkv("step", self.step + self.resume_step) 314 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 315 | if self.use_fp16: 316 | logger.logkv("lg_loss_scale", self.lg_loss_scale) 317 | 318 | def save(self): 319 | def save_checkpoint(rate, params): 320 | state_dict = self._master_params_to_state_dict(params) 321 | if dist.get_rank() == 0: 322 | logger.log(f"saving model {rate}...") 323 | if not rate: 324 | filename = f"model{(self.step+self.resume_step):06d}.pt" 325 | else: 326 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 327 | # with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 328 | with bf.BlobFile(bf.join('/root/autodl-tmp/TASK8/improved_diffusion/model_save_FULL/', filename), "wb") as f: 329 | th.save(state_dict, f) 330 | 331 | save_checkpoint(0, self.master_params) 332 | for rate, params in zip(self.ema_rate, self.ema_params): 333 | save_checkpoint(rate, params) 334 | 335 | if dist.get_rank() == 0: 336 | with bf.BlobFile( 337 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 338 | "wb", 339 | ) as f: 340 | th.save(self.opt.state_dict(), f) 341 | 342 | dist.barrier() 343 | 344 | def _master_params_to_state_dict(self, master_params): 345 | if self.use_fp16: 346 | master_params = unflatten_master_params( 347 | self.model.parameters(), master_params 348 | ) 349 | state_dict = self.model.state_dict() 350 | for i, (name, _value) in enumerate(self.model.named_parameters()): 351 | assert name in state_dict 352 | state_dict[name] = master_params[i] 353 | return state_dict 354 | 355 | def _state_dict_to_master_params(self, state_dict): 356 | params = [state_dict[name] for name, _ in self.model.named_parameters()] 357 | if self.use_fp16: 358 | return make_master_params(params) 359 | else: 360 | return params 361 | 362 | 363 | def parse_resume_step_from_filename(filename): 364 | """ 365 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 366 | checkpoint's number of steps. 367 | """ 368 | split = filename.split("model") 369 | if len(split) < 2: 370 | return 0 371 | split1 = split[-1].split(".")[0] 372 | try: 373 | return int(split1) 374 | except ValueError: 375 | return 0 376 | 377 | 378 | def get_blob_logdir(): 379 | return os.environ.get("DIFFUSION_BLOB_LOGDIR", logger.get_dir()) 380 | 381 | 382 | def find_resume_checkpoint(): 383 | # On your infrastructure, you may want to override this to automatically 384 | # discover the latest checkpoint on your blob storage, etc. 385 | return None 386 | 387 | 388 | def find_ema_checkpoint(main_checkpoint, step, rate): 389 | if main_checkpoint is None: 390 | return None 391 | filename = f"ema_{rate}_{(step):06d}.pt" 392 | path = bf.join(bf.dirname(main_checkpoint), filename) 393 | if bf.exists(path): 394 | return path 395 | return None 396 | 397 | 398 | def log_loss_dict(diffusion, ts, losses): 399 | for key, values in losses.items(): 400 | logger.logkv_mean(key, values.mean().item()) 401 | # Log the quantiles (four quartiles, in particular). 402 | 403 | # for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 404 | # quartile = int(4 * sub_t / diffusion.num_timesteps) 405 | # logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 406 | --------------------------------------------------------------------------------