├── evaluations ├── requirements.txt ├── README.md └── evaluator.py ├── guided_diffusion ├── __init__.py ├── dist_util.py ├── losses.py ├── respace.py ├── nn.py ├── resample.py ├── image_datasets.py ├── fp16_util.py ├── train_util.py ├── script_util.py ├── logger.py ├── unet.py └── gaussian_diffusion.py ├── setup.py ├── LICENSE ├── datasets ├── lsun_bedroom.py └── README.md ├── scripts ├── image_train.py ├── super_res_train.py ├── image_nll.py ├── super_res_sample.py ├── image_sample.py ├── classifier_sample.py └── classifier_train.py └── README.md /evaluations/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow-gpu>=2.0 2 | scipy 3 | requests 4 | tqdm -------------------------------------------------------------------------------- /guided_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="guided-diffusion", 5 | py_modules=["guided_diffusion"], 6 | install_requires=["blobfile>=1.0.5", "torch", "tqdm"], 7 | ) 8 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 OpenAI 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /datasets/lsun_bedroom.py: -------------------------------------------------------------------------------- 1 | """ 2 | Convert an LSUN lmdb database into a directory of images. 3 | """ 4 | 5 | import argparse 6 | import io 7 | import os 8 | 9 | from PIL import Image 10 | import lmdb 11 | import numpy as np 12 | 13 | 14 | def read_images(lmdb_path, image_size): 15 | env = lmdb.open(lmdb_path, map_size=1099511627776, max_readers=100, readonly=True) 16 | with env.begin(write=False) as transaction: 17 | cursor = transaction.cursor() 18 | for _, webp_data in cursor: 19 | img = Image.open(io.BytesIO(webp_data)) 20 | width, height = img.size 21 | scale = image_size / min(width, height) 22 | img = img.resize( 23 | (int(round(scale * width)), int(round(scale * height))), 24 | resample=Image.BOX, 25 | ) 26 | arr = np.array(img) 27 | h, w, _ = arr.shape 28 | h_off = (h - image_size) // 2 29 | w_off = (w - image_size) // 2 30 | arr = arr[h_off : h_off + image_size, w_off : w_off + image_size] 31 | yield arr 32 | 33 | 34 | def dump_images(out_dir, images, prefix): 35 | if not os.path.exists(out_dir): 36 | os.mkdir(out_dir) 37 | for i, img in enumerate(images): 38 | Image.fromarray(img).save(os.path.join(out_dir, f"{prefix}_{i:07d}.png")) 39 | 40 | 41 | def main(): 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument("--image-size", help="new image size", type=int, default=256) 44 | parser.add_argument("--prefix", help="class name", type=str, default="bedroom") 45 | parser.add_argument("lmdb_path", help="path to an LSUN lmdb database") 46 | parser.add_argument("out_dir", help="path to output directory") 47 | args = parser.parse_args() 48 | 49 | images = read_images(args.lmdb_path, args.image_size) 50 | dump_images(args.out_dir, images, args.prefix) 51 | 52 | 53 | if __name__ == "__main__": 54 | main() 55 | -------------------------------------------------------------------------------- /datasets/README.md: -------------------------------------------------------------------------------- 1 | # Downloading datasets 2 | 3 | This directory includes instructions and scripts for downloading ImageNet and LSUN bedrooms for use in this codebase. 4 | 5 | ## Class-conditional ImageNet 6 | 7 | For our class-conditional models, we use the official ILSVRC2012 dataset with manual center cropping and downsampling. To obtain this dataset, navigate to [this page on image-net.org](http://www.image-net.org/challenges/LSVRC/2012/downloads) and sign in (or create an account if you do not already have one). Then click on the link reading "Training images (Task 1 & 2)". This is a 138GB tar file containing 1000 sub-tar files, one per class. 8 | 9 | Once the file is downloaded, extract it and look inside. You should see 1000 `.tar` files. You need to extract each of these, which may be impractical to do by hand on your operating system. To automate the process on a Unix-based system, you can `cd` into the directory and run this short shell script: 10 | 11 | ``` 12 | for file in *.tar; do tar xf "$file"; rm "$file"; done 13 | ``` 14 | 15 | This will extract and remove each tar file in turn. 16 | 17 | Once all of the images have been extracted, the resulting directory should be usable as a data directory (the `--data_dir` argument for the training script). The filenames should all start with WNID (class ids) followed by underscores, like `n01440764_2708.JPEG`. Conveniently (but not by accident) this is how the automated data-loader expects to discover class labels. 18 | 19 | ## LSUN bedroom 20 | 21 | To download and pre-process LSUN bedroom, clone [fyu/lsun](https://github.com/fyu/lsun) on GitHub and run their download script `python3 download.py bedroom`. The result will be an "lmdb" database named like `bedroom_train_lmdb`. You can pass this to our [lsun_bedroom.py](lsun_bedroom.py) script like so: 22 | 23 | ``` 24 | python lsun_bedroom.py bedroom_train_lmdb lsun_train_output_dir 25 | ``` 26 | 27 | This creates a directory called `lsun_train_output_dir`. This directory can be passed to the training scripts via the `--data_dir` argument. 28 | -------------------------------------------------------------------------------- /scripts/image_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | 7 | from guided_diffusion import dist_util, logger 8 | from guided_diffusion.image_datasets import load_data 9 | from guided_diffusion.resample import create_named_schedule_sampler 10 | from guided_diffusion.script_util import ( 11 | model_and_diffusion_defaults, 12 | create_model_and_diffusion, 13 | args_to_dict, 14 | add_dict_to_argparser, 15 | ) 16 | from guided_diffusion.train_util import TrainLoop 17 | 18 | 19 | def main(): 20 | args = create_argparser().parse_args() 21 | 22 | dist_util.setup_dist() 23 | logger.configure(dir=args.log_dir) 24 | 25 | logger.log("creating model and diffusion...") 26 | model, diffusion = create_model_and_diffusion( 27 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 28 | ) 29 | model.to(dist_util.dev()) 30 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 31 | 32 | logger.log("creating data loader...") 33 | data = load_data( 34 | data_dir=args.data_dir, 35 | batch_size=args.batch_size, 36 | image_size=args.image_size, 37 | class_cond=args.class_cond, 38 | ) 39 | 40 | logger.log("training...") 41 | TrainLoop( 42 | model=model, 43 | diffusion=diffusion, 44 | data=data, 45 | batch_size=args.batch_size, 46 | microbatch=args.microbatch, 47 | lr=args.lr, 48 | ema_rate=args.ema_rate, 49 | log_interval=args.log_interval, 50 | save_interval=args.save_interval, 51 | resume_checkpoint=args.resume_checkpoint, 52 | use_fp16=args.use_fp16, 53 | fp16_scale_growth=args.fp16_scale_growth, 54 | schedule_sampler=schedule_sampler, 55 | weight_decay=args.weight_decay, 56 | lr_anneal_steps=args.lr_anneal_steps, 57 | ).run_loop() 58 | 59 | 60 | def create_argparser(): 61 | defaults = dict( 62 | data_dir="", 63 | log_dir="", 64 | schedule_sampler="uniform", 65 | lr=1e-4, 66 | weight_decay=0.0, 67 | lr_anneal_steps=0, 68 | batch_size=1, 69 | microbatch=-1, # -1 disables microbatches 70 | ema_rate="0.9999", # comma-separated list of EMA values 71 | log_interval=10, 72 | save_interval=10000, 73 | resume_checkpoint="", 74 | use_fp16=False, 75 | fp16_scale_growth=1e-3, 76 | ) 77 | defaults.update(model_and_diffusion_defaults()) 78 | parser = argparse.ArgumentParser() 79 | add_dict_to_argparser(parser, defaults) 80 | return parser 81 | 82 | 83 | if __name__ == "__main__": 84 | main() 85 | -------------------------------------------------------------------------------- /guided_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 | os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" 28 | 29 | comm = MPI.COMM_WORLD 30 | backend = "gloo" if not th.cuda.is_available() else "nccl" 31 | 32 | if backend == "gloo": 33 | hostname = "localhost" 34 | else: 35 | hostname = socket.gethostbyname(socket.getfqdn()) 36 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 37 | os.environ["RANK"] = str(comm.rank) 38 | os.environ["WORLD_SIZE"] = str(comm.size) 39 | 40 | port = comm.bcast(_find_free_port(), root=0) 41 | os.environ["MASTER_PORT"] = str(port) 42 | dist.init_process_group(backend=backend, init_method="env://") 43 | 44 | 45 | def dev(): 46 | """ 47 | Get the device to use for torch.distributed. 48 | """ 49 | if th.cuda.is_available(): 50 | return th.device(f"cuda") 51 | return th.device("cpu") 52 | 53 | 54 | def load_state_dict(path, **kwargs): 55 | """ 56 | Load a PyTorch file without redundant fetches across MPI ranks. 57 | """ 58 | chunk_size = 2 ** 30 # MPI has a relatively small size limit 59 | if MPI.COMM_WORLD.Get_rank() == 0: 60 | with bf.BlobFile(path, "rb") as f: 61 | data = f.read() 62 | num_chunks = len(data) // chunk_size 63 | if len(data) % chunk_size: 64 | num_chunks += 1 65 | MPI.COMM_WORLD.bcast(num_chunks) 66 | for i in range(0, len(data), chunk_size): 67 | MPI.COMM_WORLD.bcast(data[i : i + chunk_size]) 68 | else: 69 | num_chunks = MPI.COMM_WORLD.bcast(None) 70 | data = bytes() 71 | for _ in range(num_chunks): 72 | data += MPI.COMM_WORLD.bcast(None) 73 | 74 | return th.load(io.BytesIO(data), **kwargs) 75 | 76 | 77 | def sync_params(params): 78 | """ 79 | Synchronize a sequence of Tensors across ranks from rank 0. 80 | """ 81 | for p in params: 82 | with th.no_grad(): 83 | dist.broadcast(p, 0) 84 | 85 | 86 | def _find_free_port(): 87 | try: 88 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 89 | s.bind(("", 0)) 90 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 91 | return s.getsockname()[1] 92 | finally: 93 | s.close() 94 | -------------------------------------------------------------------------------- /guided_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 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # P2 weighting (CVPR 2022) 2 | 3 | This is the codebase for [Perception Prioritized Training of Diffusion Models](https://arxiv.org/abs/2204.00227). 4 | 5 | This repository is heavily based on [openai/guided-diffusion](https://github.com/openai/guided-diffusion). 6 | 7 | P2 modifies the weighting scheme of the training objective function to improve sample quality. It encourages the diffusion model to focus on recovering signals from highly corrupted data, where the model learns global and perceptually rich concepts. Below figure shows the weighting schemes in terms of SNR. 8 | 9 | ![snr_weight](https://user-images.githubusercontent.com/36615789/161203299-8b02d76b-9c51-4529-8329-3ac08e9f3bc8.png) 10 | 11 | ## Pre-trained models 12 | 13 | All models are trained at 256x256 resolution. 14 | 15 | Here are the models trained on FFHQ, CelebA-HQ, CUB, AFHQ-Dogs, Flowers, and MetFaces: [onedrive](https://1drv.ms/f/s!AkQjJhxDm0Fyhqp_4gkYjwVRBe8V_w?e=Us79E9) [gdrive](https://drive.google.com/drive/folders/1bcWh3XuQzdct4-UPTrIX-lvs47OiLaOM?usp=sharing) 16 | 17 | ## Requirements 18 | We tested on PyTorch 1.7.1, single RTX8000 GPU. 19 | 20 | ## Sampling from pre-trained models 21 | 22 | First, set PYTHONPATH variable to point to the root of the repository. Do the same when training new models. 23 | 24 | ``` 25 | export PYTHONPATH=$PYTHONPATH:$(pwd) 26 | ``` 27 | 28 | Put model checkpoints into a folder `models/`. 29 | 30 | Samples will be saved in `samples/`. 31 | 32 | ``` 33 | python scripts/image_sample.py --attention_resolutions 16 --class_cond False --diffusion_steps 1000 --dropout 0.0 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 128 --num_res_blocks 1 --num_head_channels 64 --resblock_updown True --use_fp16 False --use_scale_shift_norm True --timestep_respacing 250 --model_path models/ffhq_p2.pt --sample_dir samples 34 | ``` 35 | 36 | To sample for 250 timesteps without DDIM, replace `--timestep_respacing ddim25` to `--timestep_respacing 250`, and replace `--use_ddim True` with `--use_ddim False`. 37 | 38 | ## Training your models 39 | 40 | `--p2_gamma` and `--p2_k` are two hyperparameters of P2 weighting. We used `--p2_gamma 0.5 --p2_k 1` and `--p2_gamma 1 --p2_k 1` in the paper. 41 | 42 | Logs and models will be saved in `logs/`. You should modify `--data_dir`. 43 | 44 | We used lightweight version (93M parameter) of [ADM](https://arxiv.org/abs/2105.05233) (over 500M) as default model configuration. You may modify the model. 45 | 46 | ``` 47 | python scripts/image_train.py --data_dir data/DATASET_NAME --attention_resolutions 16 --class_cond False --diffusion_steps 1000 --dropout 0.0 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 128 --num_head_channels 64 --num_res_blocks 1 --resblock_updown True --use_fp16 False --use_scale_shift_norm True --lr 2e-5 --batch_size 8 --rescale_learned_sigmas True --p2_gamma 1 --p2_k 1 --log_dir logs 48 | ``` 49 | 50 | 51 | -------------------------------------------------------------------------------- /scripts/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 | 9 | from guided_diffusion import dist_util, logger 10 | from guided_diffusion.image_datasets import load_data 11 | from guided_diffusion.resample import create_named_schedule_sampler 12 | from guided_diffusion.script_util import ( 13 | sr_model_and_diffusion_defaults, 14 | sr_create_model_and_diffusion, 15 | args_to_dict, 16 | add_dict_to_argparser, 17 | ) 18 | from guided_diffusion.train_util import TrainLoop 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...") 28 | model, diffusion = sr_create_model_and_diffusion( 29 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 30 | ) 31 | model.to(dist_util.dev()) 32 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 33 | 34 | logger.log("creating data loader...") 35 | data = load_superres_data( 36 | args.data_dir, 37 | args.batch_size, 38 | large_size=args.large_size, 39 | small_size=args.small_size, 40 | class_cond=args.class_cond, 41 | ) 42 | 43 | logger.log("training...") 44 | TrainLoop( 45 | model=model, 46 | diffusion=diffusion, 47 | data=data, 48 | batch_size=args.batch_size, 49 | microbatch=args.microbatch, 50 | lr=args.lr, 51 | ema_rate=args.ema_rate, 52 | log_interval=args.log_interval, 53 | save_interval=args.save_interval, 54 | resume_checkpoint=args.resume_checkpoint, 55 | use_fp16=args.use_fp16, 56 | fp16_scale_growth=args.fp16_scale_growth, 57 | schedule_sampler=schedule_sampler, 58 | weight_decay=args.weight_decay, 59 | lr_anneal_steps=args.lr_anneal_steps, 60 | ).run_loop() 61 | 62 | 63 | def load_superres_data(data_dir, batch_size, large_size, small_size, class_cond=False): 64 | data = load_data( 65 | data_dir=data_dir, 66 | batch_size=batch_size, 67 | image_size=large_size, 68 | class_cond=class_cond, 69 | ) 70 | for large_batch, model_kwargs in data: 71 | model_kwargs["low_res"] = F.interpolate(large_batch, small_size, mode="area") 72 | yield large_batch, model_kwargs 73 | 74 | 75 | def create_argparser(): 76 | defaults = dict( 77 | data_dir="", 78 | schedule_sampler="uniform", 79 | lr=1e-4, 80 | weight_decay=0.0, 81 | lr_anneal_steps=0, 82 | batch_size=1, 83 | microbatch=-1, 84 | ema_rate="0.9999", 85 | log_interval=10, 86 | save_interval=10000, 87 | resume_checkpoint="", 88 | use_fp16=False, 89 | fp16_scale_growth=1e-3, 90 | ) 91 | defaults.update(sr_model_and_diffusion_defaults()) 92 | parser = argparse.ArgumentParser() 93 | add_dict_to_argparser(parser, defaults) 94 | return parser 95 | 96 | 97 | if __name__ == "__main__": 98 | main() 99 | -------------------------------------------------------------------------------- /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 guided_diffusion import dist_util, logger 12 | from guided_diffusion.image_datasets import load_data 13 | from guided_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 | -------------------------------------------------------------------------------- /scripts/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 | 6 | import argparse 7 | import os 8 | 9 | import blobfile as bf 10 | import numpy as np 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | from guided_diffusion import dist_util, logger 15 | from guided_diffusion.script_util import ( 16 | sr_model_and_diffusion_defaults, 17 | sr_create_model_and_diffusion, 18 | args_to_dict, 19 | add_dict_to_argparser, 20 | ) 21 | 22 | 23 | def main(): 24 | args = create_argparser().parse_args() 25 | 26 | dist_util.setup_dist() 27 | logger.configure() 28 | 29 | logger.log("creating model...") 30 | model, diffusion = sr_create_model_and_diffusion( 31 | **args_to_dict(args, sr_model_and_diffusion_defaults().keys()) 32 | ) 33 | model.load_state_dict( 34 | dist_util.load_state_dict(args.model_path, map_location="cpu") 35 | ) 36 | model.to(dist_util.dev()) 37 | if args.use_fp16: 38 | model.convert_to_fp16() 39 | model.eval() 40 | 41 | logger.log("loading data...") 42 | data = load_data_for_worker(args.base_samples, args.batch_size, args.class_cond) 43 | 44 | logger.log("creating samples...") 45 | all_images = [] 46 | while len(all_images) * args.batch_size < args.num_samples: 47 | model_kwargs = next(data) 48 | model_kwargs = {k: v.to(dist_util.dev()) for k, v in model_kwargs.items()} 49 | sample = diffusion.p_sample_loop( 50 | model, 51 | (args.batch_size, 3, args.large_size, args.large_size), 52 | clip_denoised=args.clip_denoised, 53 | model_kwargs=model_kwargs, 54 | ) 55 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 56 | sample = sample.permute(0, 2, 3, 1) 57 | sample = sample.contiguous() 58 | 59 | all_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 60 | dist.all_gather(all_samples, sample) # gather not supported with NCCL 61 | for sample in all_samples: 62 | all_images.append(sample.cpu().numpy()) 63 | logger.log(f"created {len(all_images) * args.batch_size} samples") 64 | 65 | arr = np.concatenate(all_images, axis=0) 66 | arr = arr[: args.num_samples] 67 | if dist.get_rank() == 0: 68 | shape_str = "x".join([str(x) for x in arr.shape]) 69 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 70 | logger.log(f"saving to {out_path}") 71 | np.savez(out_path, arr) 72 | 73 | dist.barrier() 74 | logger.log("sampling complete") 75 | 76 | 77 | def load_data_for_worker(base_samples, batch_size, class_cond): 78 | with bf.BlobFile(base_samples, "rb") as f: 79 | obj = np.load(f) 80 | image_arr = obj["arr_0"] 81 | if class_cond: 82 | label_arr = obj["arr_1"] 83 | rank = dist.get_rank() 84 | num_ranks = dist.get_world_size() 85 | buffer = [] 86 | label_buffer = [] 87 | while True: 88 | for i in range(rank, len(image_arr), num_ranks): 89 | buffer.append(image_arr[i]) 90 | if class_cond: 91 | label_buffer.append(label_arr[i]) 92 | if len(buffer) == batch_size: 93 | batch = th.from_numpy(np.stack(buffer)).float() 94 | batch = batch / 127.5 - 1.0 95 | batch = batch.permute(0, 3, 1, 2) 96 | res = dict(low_res=batch) 97 | if class_cond: 98 | res["y"] = th.from_numpy(np.stack(label_buffer)) 99 | yield res 100 | buffer, label_buffer = [], [] 101 | 102 | 103 | def create_argparser(): 104 | defaults = dict( 105 | clip_denoised=True, 106 | num_samples=10000, 107 | batch_size=16, 108 | use_ddim=False, 109 | base_samples="", 110 | model_path="", 111 | ) 112 | defaults.update(sr_model_and_diffusion_defaults()) 113 | parser = argparse.ArgumentParser() 114 | add_dict_to_argparser(parser, defaults) 115 | return parser 116 | 117 | 118 | if __name__ == "__main__": 119 | main() 120 | -------------------------------------------------------------------------------- /scripts/image_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of image samples from a model and save them as a large 3 | numpy array. This can be used to produce samples for FID evaluation. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | 13 | from guided_diffusion import dist_util, logger 14 | from guided_diffusion.script_util import ( 15 | NUM_CLASSES, 16 | model_and_diffusion_defaults, 17 | create_model_and_diffusion, 18 | add_dict_to_argparser, 19 | args_to_dict, 20 | ) 21 | from torchvision import utils 22 | 23 | 24 | def main(): 25 | args = create_argparser().parse_args() 26 | 27 | dist_util.setup_dist() 28 | logger.configure(dir=args.sample_dir) 29 | 30 | logger.log("creating model and diffusion...") 31 | model, diffusion = create_model_and_diffusion( 32 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 33 | ) 34 | model.load_state_dict( 35 | dist_util.load_state_dict(args.model_path, map_location="cpu") 36 | ) 37 | model.to(dist_util.dev()) 38 | if args.use_fp16: 39 | model.convert_to_fp16() 40 | model.eval() 41 | 42 | logger.log("sampling...") 43 | all_images = [] 44 | all_labels = [] 45 | count = 0 46 | while count * args.batch_size < args.num_samples: 47 | model_kwargs = {} 48 | if args.class_cond: 49 | classes = th.randint( 50 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 51 | ) 52 | model_kwargs["y"] = classes 53 | sample_fn = ( 54 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 55 | ) 56 | sample = sample_fn( 57 | model, 58 | (args.batch_size, 3, args.image_size, args.image_size), 59 | clip_denoised=args.clip_denoised, 60 | model_kwargs=model_kwargs, 61 | ) 62 | # saving png 63 | for i in range(args.batch_size): 64 | out_path = os.path.join(logger.get_dir(), 65 | f"{str(count * args.batch_size + i).zfill(5)}.png") 66 | utils.save_image( 67 | sample[i].unsqueeze(0), 68 | out_path, 69 | nrow=1, 70 | normalize=True, 71 | range=(-1, 1), 72 | ) 73 | # saving npz 74 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 75 | sample = sample.permute(0, 2, 3, 1) 76 | sample = sample.contiguous() 77 | 78 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 79 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 80 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 81 | if args.class_cond: 82 | gathered_labels = [ 83 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 84 | ] 85 | dist.all_gather(gathered_labels, classes) 86 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 87 | logger.log(f"created {len(all_images) * args.batch_size} samples") 88 | 89 | arr = np.concatenate(all_images, axis=0) 90 | arr = arr[: args.num_samples] 91 | if args.class_cond: 92 | label_arr = np.concatenate(all_labels, axis=0) 93 | label_arr = label_arr[: args.num_samples] 94 | if dist.get_rank() == 0: 95 | shape_str = "x".join([str(x) for x in arr.shape]) 96 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 97 | logger.log(f"saving to {out_path}") 98 | if args.class_cond: 99 | np.savez(out_path, arr, label_arr) 100 | else: 101 | np.savez(out_path, arr) 102 | 103 | dist.barrier() 104 | logger.log("sampling complete") 105 | 106 | 107 | def create_argparser(): 108 | defaults = dict( 109 | clip_denoised=True, 110 | num_samples=10000, 111 | batch_size=16, 112 | use_ddim=False, 113 | model_path="", 114 | sample_dir="", 115 | ) 116 | defaults.update(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 | -------------------------------------------------------------------------------- /scripts/classifier_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Like image_sample.py, but use a noisy image classifier to guide the sampling 3 | process towards more realistic images. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | import torch.nn.functional as F 13 | 14 | from guided_diffusion import dist_util, logger 15 | from guided_diffusion.script_util import ( 16 | NUM_CLASSES, 17 | model_and_diffusion_defaults, 18 | classifier_defaults, 19 | create_model_and_diffusion, 20 | create_classifier, 21 | add_dict_to_argparser, 22 | args_to_dict, 23 | ) 24 | 25 | 26 | def main(): 27 | args = create_argparser().parse_args() 28 | 29 | dist_util.setup_dist() 30 | logger.configure() 31 | 32 | logger.log("creating model and diffusion...") 33 | model, diffusion = create_model_and_diffusion( 34 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 35 | ) 36 | model.load_state_dict( 37 | dist_util.load_state_dict(args.model_path, map_location="cpu") 38 | ) 39 | model.to(dist_util.dev()) 40 | if args.use_fp16: 41 | model.convert_to_fp16() 42 | model.eval() 43 | 44 | logger.log("loading classifier...") 45 | classifier = create_classifier(**args_to_dict(args, classifier_defaults().keys())) 46 | classifier.load_state_dict( 47 | dist_util.load_state_dict(args.classifier_path, map_location="cpu") 48 | ) 49 | classifier.to(dist_util.dev()) 50 | if args.classifier_use_fp16: 51 | classifier.convert_to_fp16() 52 | classifier.eval() 53 | 54 | def cond_fn(x, t, y=None): 55 | assert y is not None 56 | with th.enable_grad(): 57 | x_in = x.detach().requires_grad_(True) 58 | logits = classifier(x_in, t) 59 | log_probs = F.log_softmax(logits, dim=-1) 60 | selected = log_probs[range(len(logits)), y.view(-1)] 61 | return th.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale 62 | 63 | def model_fn(x, t, y=None): 64 | assert y is not None 65 | return model(x, t, y if args.class_cond else None) 66 | 67 | logger.log("sampling...") 68 | all_images = [] 69 | all_labels = [] 70 | while len(all_images) * args.batch_size < args.num_samples: 71 | model_kwargs = {} 72 | classes = th.randint( 73 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 74 | ) 75 | model_kwargs["y"] = classes 76 | sample_fn = ( 77 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 78 | ) 79 | sample = sample_fn( 80 | model_fn, 81 | (args.batch_size, 3, args.image_size, args.image_size), 82 | clip_denoised=args.clip_denoised, 83 | model_kwargs=model_kwargs, 84 | cond_fn=cond_fn, 85 | device=dist_util.dev(), 86 | ) 87 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 88 | sample = sample.permute(0, 2, 3, 1) 89 | sample = sample.contiguous() 90 | 91 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 92 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 93 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 94 | gathered_labels = [th.zeros_like(classes) for _ in range(dist.get_world_size())] 95 | dist.all_gather(gathered_labels, classes) 96 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 97 | logger.log(f"created {len(all_images) * args.batch_size} samples") 98 | 99 | arr = np.concatenate(all_images, axis=0) 100 | arr = arr[: args.num_samples] 101 | label_arr = np.concatenate(all_labels, axis=0) 102 | label_arr = label_arr[: args.num_samples] 103 | if dist.get_rank() == 0: 104 | shape_str = "x".join([str(x) for x in arr.shape]) 105 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 106 | logger.log(f"saving to {out_path}") 107 | np.savez(out_path, arr, label_arr) 108 | 109 | dist.barrier() 110 | logger.log("sampling complete") 111 | 112 | 113 | def create_argparser(): 114 | defaults = dict( 115 | clip_denoised=True, 116 | num_samples=10000, 117 | batch_size=16, 118 | use_ddim=False, 119 | model_path="", 120 | classifier_path="", 121 | classifier_scale=1.0, 122 | ) 123 | defaults.update(model_and_diffusion_defaults()) 124 | defaults.update(classifier_defaults()) 125 | parser = argparse.ArgumentParser() 126 | add_dict_to_argparser(parser, defaults) 127 | return parser 128 | 129 | 130 | if __name__ == "__main__": 131 | main() 132 | -------------------------------------------------------------------------------- /guided_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 condition_mean(self, cond_fn, *args, **kwargs): 99 | return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) 100 | 101 | def condition_score(self, cond_fn, *args, **kwargs): 102 | return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) 103 | 104 | def _wrap_model(self, model): 105 | if isinstance(model, _WrappedModel): 106 | return model 107 | return _WrappedModel( 108 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps 109 | ) 110 | 111 | def _scale_timesteps(self, t): 112 | # Scaling is done by the wrapped model. 113 | return t 114 | 115 | 116 | class _WrappedModel: 117 | def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): 118 | self.model = model 119 | self.timestep_map = timestep_map 120 | self.rescale_timesteps = rescale_timesteps 121 | self.original_num_steps = original_num_steps 122 | 123 | def __call__(self, x, ts, **kwargs): 124 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 125 | new_ts = map_tensor[ts] 126 | if self.rescale_timesteps: 127 | new_ts = new_ts.float() * (1000.0 / self.original_num_steps) 128 | return self.model(x, new_ts, **kwargs) 129 | -------------------------------------------------------------------------------- /guided_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 | -------------------------------------------------------------------------------- /guided_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 | 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 | -------------------------------------------------------------------------------- /guided_diffusion/image_datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | 4 | from PIL import Image 5 | import blobfile as bf 6 | from mpi4py import MPI 7 | import numpy as np 8 | from torch.utils.data import DataLoader, Dataset 9 | 10 | 11 | def load_data( 12 | *, 13 | data_dir, 14 | batch_size, 15 | image_size, 16 | class_cond=False, 17 | deterministic=False, 18 | random_crop=False, 19 | random_flip=True, 20 | ): 21 | """ 22 | For a dataset, create a generator over (images, kwargs) pairs. 23 | 24 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 25 | more keys, each of which map to a batched Tensor of their own. 26 | The kwargs dict can be used for class labels, in which case the key is "y" 27 | and the values are integer tensors of class labels. 28 | 29 | :param data_dir: a dataset directory. 30 | :param batch_size: the batch size of each returned pair. 31 | :param image_size: the size to which images are resized. 32 | :param class_cond: if True, include a "y" key in returned dicts for class 33 | label. If classes are not available and this is true, an 34 | exception will be raised. 35 | :param deterministic: if True, yield results in a deterministic order. 36 | :param random_crop: if True, randomly crop the images for augmentation. 37 | :param random_flip: if True, randomly flip the images for augmentation. 38 | """ 39 | if not data_dir: 40 | raise ValueError("unspecified data directory") 41 | all_files = _list_image_files_recursively(data_dir) 42 | classes = None 43 | if class_cond: 44 | # Assume classes are the first part of the filename, 45 | # before an underscore. 46 | class_names = [bf.basename(path).split("_")[0] for path in all_files] 47 | sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 48 | classes = [sorted_classes[x] for x in class_names] 49 | dataset = ImageDataset( 50 | image_size, 51 | all_files, 52 | classes=classes, 53 | shard=MPI.COMM_WORLD.Get_rank(), 54 | num_shards=MPI.COMM_WORLD.Get_size(), 55 | random_crop=random_crop, 56 | random_flip=random_flip, 57 | ) 58 | if deterministic: 59 | loader = DataLoader( 60 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 61 | ) 62 | else: 63 | loader = DataLoader( 64 | dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True 65 | ) 66 | while True: 67 | yield from loader 68 | 69 | 70 | def _list_image_files_recursively(data_dir): 71 | results = [] 72 | for entry in sorted(bf.listdir(data_dir)): 73 | full_path = bf.join(data_dir, entry) 74 | ext = entry.split(".")[-1] 75 | if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]: 76 | results.append(full_path) 77 | elif bf.isdir(full_path): 78 | results.extend(_list_image_files_recursively(full_path)) 79 | return results 80 | 81 | 82 | class ImageDataset(Dataset): 83 | def __init__( 84 | self, 85 | resolution, 86 | image_paths, 87 | classes=None, 88 | shard=0, 89 | num_shards=1, 90 | random_crop=False, 91 | random_flip=True, 92 | ): 93 | super().__init__() 94 | self.resolution = resolution 95 | self.local_images = image_paths[shard:][::num_shards] 96 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 97 | self.random_crop = random_crop 98 | self.random_flip = random_flip 99 | 100 | def __len__(self): 101 | return len(self.local_images) 102 | 103 | def __getitem__(self, idx): 104 | path = self.local_images[idx] 105 | with bf.BlobFile(path, "rb") as f: 106 | pil_image = Image.open(f) 107 | pil_image.load() 108 | pil_image = pil_image.convert("RGB") 109 | 110 | if self.random_crop: 111 | arr = random_crop_arr(pil_image, self.resolution) 112 | else: 113 | arr = center_crop_arr(pil_image, self.resolution) 114 | 115 | if self.random_flip and random.random() < 0.5: 116 | arr = arr[:, ::-1] 117 | 118 | arr = arr.astype(np.float32) / 127.5 - 1 119 | 120 | out_dict = {} 121 | if self.local_classes is not None: 122 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 123 | return np.transpose(arr, [2, 0, 1]), out_dict 124 | 125 | 126 | def center_crop_arr(pil_image, image_size): 127 | # We are not on a new enough PIL to support the `reducing_gap` 128 | # argument, which uses BOX downsampling at powers of two first. 129 | # Thus, we do it by hand to improve downsample quality. 130 | while min(*pil_image.size) >= 2 * image_size: 131 | pil_image = pil_image.resize( 132 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 133 | ) 134 | 135 | scale = image_size / min(*pil_image.size) 136 | pil_image = pil_image.resize( 137 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 138 | ) 139 | 140 | arr = np.array(pil_image) 141 | crop_y = (arr.shape[0] - image_size) // 2 142 | crop_x = (arr.shape[1] - image_size) // 2 143 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 144 | 145 | 146 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 147 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 148 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 149 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 150 | 151 | # We are not on a new enough PIL to support the `reducing_gap` 152 | # argument, which uses BOX downsampling at powers of two first. 153 | # Thus, we do it by hand to improve downsample quality. 154 | while min(*pil_image.size) >= 2 * smaller_dim_size: 155 | pil_image = pil_image.resize( 156 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 157 | ) 158 | 159 | scale = smaller_dim_size / min(*pil_image.size) 160 | pil_image = pil_image.resize( 161 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 162 | ) 163 | 164 | arr = np.array(pil_image) 165 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 166 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 167 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 168 | -------------------------------------------------------------------------------- /evaluations/README.md: -------------------------------------------------------------------------------- 1 | # Evaluations 2 | 3 | To compare different generative models, we use FID, sFID, Precision, Recall, and Inception Score. These metrics can all be calculated using batches of samples, which we store in `.npz` (numpy) files. 4 | 5 | # Download batches 6 | 7 | We provide pre-computed sample batches for the reference datasets, our diffusion models, and several baselines we compare against. These are all stored in `.npz` format. 8 | 9 | Reference dataset batches contain pre-computed statistics over the whole dataset, as well as 10,000 images for computing Precision and Recall. All other batches contain 50,000 images which can be used to compute statistics and Precision/Recall. 10 | 11 | Here are links to download all of the sample and reference batches: 12 | 13 | * LSUN 14 | * LSUN bedroom: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/VIRTUAL_lsun_bedroom256.npz) 15 | * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/admnet_dropout_lsun_bedroom.npz) 16 | * [DDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/ddpm_lsun_bedroom.npz) 17 | * [IDDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/iddpm_lsun_bedroom.npz) 18 | * [StyleGAN](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/bedroom/stylegan_lsun_bedroom.npz) 19 | * LSUN cat: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/VIRTUAL_lsun_cat256.npz) 20 | * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/admnet_dropout_lsun_cat.npz) 21 | * [StyleGAN2](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/cat/stylegan2_lsun_cat.npz) 22 | * LSUN horse: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/VIRTUAL_lsun_horse256.npz) 23 | * [ADM (dropout)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/admnet_dropout_lsun_horse.npz) 24 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/lsun/horse/admnet_lsun_horse.npz) 25 | * ImageNet 64x64: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/VIRTUAL_imagenet64_labeled.npz) 26 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/admnet_imagenet64.npz) 27 | * [IDDPM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/iddpm_imagenet64.npz) 28 | * [BigGAN](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/biggan_deep_imagenet64.npz) 29 | * ImageNet 128x128: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/VIRTUAL_imagenet128_labeled.npz) 30 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_imagenet128.npz) 31 | * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_guided_imagenet128.npz) 32 | * [ADM-G, 25 steps](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/admnet_guided_25step_imagenet128.npz) 33 | * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/biggan_deep_trunc1_imagenet128.npz) 34 | * ImageNet 256x256: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz) 35 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_imagenet256.npz) 36 | * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_imagenet256.npz) 37 | * [ADM-G, 25 step](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_25step_imagenet256.npz) 38 | * [ADM-G + ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_guided_upsampled_imagenet256.npz) 39 | * [ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/admnet_upsampled_imagenet256.npz) 40 | * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/biggan_deep_trunc1_imagenet256.npz) 41 | * ImageNet 512x512: [reference batch](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/VIRTUAL_imagenet512.npz) 42 | * [ADM](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_imagenet512.npz) 43 | * [ADM-G](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_imagenet512.npz) 44 | * [ADM-G, 25 step](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_25step_imagenet512.npz) 45 | * [ADM-G + ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_guided_upsampled_imagenet512.npz) 46 | * [ADM-U](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/admnet_upsampled_imagenet512.npz) 47 | * [BigGAN-deep (trunc=1.0)](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/biggan_deep_trunc1_imagenet512.npz) 48 | 49 | # Run evaluations 50 | 51 | First, generate or download a batch of samples and download the corresponding reference batch for the given dataset. For this example, we'll use ImageNet 256x256, so the refernce batch is `VIRTUAL_imagenet256_labeled.npz` and we can use the sample batch `admnet_guided_upsampled_imagenet256.npz`. 52 | 53 | Next, run the `evaluator.py` script. The requirements of this script can be found in [requirements.txt](requirements.txt). Pass two arguments to the script: the reference batch and the sample batch. The script will download the InceptionV3 model used for evaluations into the current working directory (if it is not already present). This file is roughly 100MB. 54 | 55 | The output of the script will look something like this, where the first `...` is a bunch of verbose TensorFlow logging: 56 | 57 | ``` 58 | $ python evaluator.py VIRTUAL_imagenet256_labeled.npz admnet_guided_upsampled_imagenet256.npz 59 | ... 60 | computing reference batch activations... 61 | computing/reading reference batch statistics... 62 | computing sample batch activations... 63 | computing/reading sample batch statistics... 64 | Computing evaluations... 65 | Inception Score: 215.8370361328125 66 | FID: 3.9425574129223264 67 | sFID: 6.140433703346162 68 | Precision: 0.8265 69 | Recall: 0.5309 70 | ``` 71 | -------------------------------------------------------------------------------- /scripts/classifier_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a noised image classifier on ImageNet. 3 | """ 4 | 5 | import argparse 6 | import os 7 | 8 | import blobfile as bf 9 | import torch as th 10 | import torch.distributed as dist 11 | import torch.nn.functional as F 12 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 13 | from torch.optim import AdamW 14 | 15 | from guided_diffusion import dist_util, logger 16 | from guided_diffusion.fp16_util import MixedPrecisionTrainer 17 | from guided_diffusion.image_datasets import load_data 18 | from guided_diffusion.resample import create_named_schedule_sampler 19 | from guided_diffusion.script_util import ( 20 | add_dict_to_argparser, 21 | args_to_dict, 22 | classifier_and_diffusion_defaults, 23 | create_classifier_and_diffusion, 24 | ) 25 | from guided_diffusion.train_util import parse_resume_step_from_filename, log_loss_dict 26 | 27 | 28 | def main(): 29 | args = create_argparser().parse_args() 30 | 31 | dist_util.setup_dist() 32 | logger.configure() 33 | 34 | logger.log("creating model and diffusion...") 35 | model, diffusion = create_classifier_and_diffusion( 36 | **args_to_dict(args, classifier_and_diffusion_defaults().keys()) 37 | ) 38 | model.to(dist_util.dev()) 39 | if args.noised: 40 | schedule_sampler = create_named_schedule_sampler( 41 | args.schedule_sampler, diffusion 42 | ) 43 | 44 | resume_step = 0 45 | if args.resume_checkpoint: 46 | resume_step = parse_resume_step_from_filename(args.resume_checkpoint) 47 | if dist.get_rank() == 0: 48 | logger.log( 49 | f"loading model from checkpoint: {args.resume_checkpoint}... at {resume_step} step" 50 | ) 51 | model.load_state_dict( 52 | dist_util.load_state_dict( 53 | args.resume_checkpoint, map_location=dist_util.dev() 54 | ) 55 | ) 56 | 57 | # Needed for creating correct EMAs and fp16 parameters. 58 | dist_util.sync_params(model.parameters()) 59 | 60 | mp_trainer = MixedPrecisionTrainer( 61 | model=model, use_fp16=args.classifier_use_fp16, initial_lg_loss_scale=16.0 62 | ) 63 | 64 | model = DDP( 65 | model, 66 | device_ids=[dist_util.dev()], 67 | output_device=dist_util.dev(), 68 | broadcast_buffers=False, 69 | bucket_cap_mb=128, 70 | find_unused_parameters=False, 71 | ) 72 | 73 | logger.log("creating data loader...") 74 | data = load_data( 75 | data_dir=args.data_dir, 76 | batch_size=args.batch_size, 77 | image_size=args.image_size, 78 | class_cond=True, 79 | random_crop=True, 80 | ) 81 | if args.val_data_dir: 82 | val_data = load_data( 83 | data_dir=args.val_data_dir, 84 | batch_size=args.batch_size, 85 | image_size=args.image_size, 86 | class_cond=True, 87 | ) 88 | else: 89 | val_data = None 90 | 91 | logger.log(f"creating optimizer...") 92 | opt = AdamW(mp_trainer.master_params, lr=args.lr, weight_decay=args.weight_decay) 93 | if args.resume_checkpoint: 94 | opt_checkpoint = bf.join( 95 | bf.dirname(args.resume_checkpoint), f"opt{resume_step:06}.pt" 96 | ) 97 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 98 | opt.load_state_dict( 99 | dist_util.load_state_dict(opt_checkpoint, map_location=dist_util.dev()) 100 | ) 101 | 102 | logger.log("training classifier model...") 103 | 104 | def forward_backward_log(data_loader, prefix="train"): 105 | batch, extra = next(data_loader) 106 | labels = extra["y"].to(dist_util.dev()) 107 | 108 | batch = batch.to(dist_util.dev()) 109 | # Noisy images 110 | if args.noised: 111 | t, _ = schedule_sampler.sample(batch.shape[0], dist_util.dev()) 112 | batch = diffusion.q_sample(batch, t) 113 | else: 114 | t = th.zeros(batch.shape[0], dtype=th.long, device=dist_util.dev()) 115 | 116 | for i, (sub_batch, sub_labels, sub_t) in enumerate( 117 | split_microbatches(args.microbatch, batch, labels, t) 118 | ): 119 | logits = model(sub_batch, timesteps=sub_t) 120 | loss = F.cross_entropy(logits, sub_labels, reduction="none") 121 | 122 | losses = {} 123 | losses[f"{prefix}_loss"] = loss.detach() 124 | losses[f"{prefix}_acc@1"] = compute_top_k( 125 | logits, sub_labels, k=1, reduction="none" 126 | ) 127 | losses[f"{prefix}_acc@5"] = compute_top_k( 128 | logits, sub_labels, k=5, reduction="none" 129 | ) 130 | log_loss_dict(diffusion, sub_t, losses) 131 | del losses 132 | loss = loss.mean() 133 | if loss.requires_grad: 134 | if i == 0: 135 | mp_trainer.zero_grad() 136 | mp_trainer.backward(loss * len(sub_batch) / len(batch)) 137 | 138 | for step in range(args.iterations - resume_step): 139 | logger.logkv("step", step + resume_step) 140 | logger.logkv( 141 | "samples", 142 | (step + resume_step + 1) * args.batch_size * dist.get_world_size(), 143 | ) 144 | if args.anneal_lr: 145 | set_annealed_lr(opt, args.lr, (step + resume_step) / args.iterations) 146 | forward_backward_log(data) 147 | mp_trainer.optimize(opt) 148 | if val_data is not None and not step % args.eval_interval: 149 | with th.no_grad(): 150 | with model.no_sync(): 151 | model.eval() 152 | forward_backward_log(val_data, prefix="val") 153 | model.train() 154 | if not step % args.log_interval: 155 | logger.dumpkvs() 156 | if ( 157 | step 158 | and dist.get_rank() == 0 159 | and not (step + resume_step) % args.save_interval 160 | ): 161 | logger.log("saving model...") 162 | save_model(mp_trainer, opt, step + resume_step) 163 | 164 | if dist.get_rank() == 0: 165 | logger.log("saving model...") 166 | save_model(mp_trainer, opt, step + resume_step) 167 | dist.barrier() 168 | 169 | 170 | def set_annealed_lr(opt, base_lr, frac_done): 171 | lr = base_lr * (1 - frac_done) 172 | for param_group in opt.param_groups: 173 | param_group["lr"] = lr 174 | 175 | 176 | def save_model(mp_trainer, opt, step): 177 | if dist.get_rank() == 0: 178 | th.save( 179 | mp_trainer.master_params_to_state_dict(mp_trainer.master_params), 180 | os.path.join(logger.get_dir(), f"model{step:06d}.pt"), 181 | ) 182 | th.save(opt.state_dict(), os.path.join(logger.get_dir(), f"opt{step:06d}.pt")) 183 | 184 | 185 | def compute_top_k(logits, labels, k, reduction="mean"): 186 | _, top_ks = th.topk(logits, k, dim=-1) 187 | if reduction == "mean": 188 | return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item() 189 | elif reduction == "none": 190 | return (top_ks == labels[:, None]).float().sum(dim=-1) 191 | 192 | 193 | def split_microbatches(microbatch, *args): 194 | bs = len(args[0]) 195 | if microbatch == -1 or microbatch >= bs: 196 | yield tuple(args) 197 | else: 198 | for i in range(0, bs, microbatch): 199 | yield tuple(x[i : i + microbatch] if x is not None else None for x in args) 200 | 201 | 202 | def create_argparser(): 203 | defaults = dict( 204 | data_dir="", 205 | val_data_dir="", 206 | noised=True, 207 | iterations=150000, 208 | lr=3e-4, 209 | weight_decay=0.0, 210 | anneal_lr=False, 211 | batch_size=4, 212 | microbatch=-1, 213 | schedule_sampler="uniform", 214 | resume_checkpoint="", 215 | log_interval=10, 216 | eval_interval=5, 217 | save_interval=10000, 218 | ) 219 | defaults.update(classifier_and_diffusion_defaults()) 220 | parser = argparse.ArgumentParser() 221 | add_dict_to_argparser(parser, defaults) 222 | return parser 223 | 224 | 225 | if __name__ == "__main__": 226 | main() 227 | -------------------------------------------------------------------------------- /guided_diffusion/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import numpy as np 6 | import torch as th 7 | import torch.nn as nn 8 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 9 | 10 | from . import logger 11 | 12 | INITIAL_LOG_LOSS_SCALE = 20.0 13 | 14 | 15 | def convert_module_to_f16(l): 16 | """ 17 | Convert primitive modules to float16. 18 | """ 19 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 20 | l.weight.data = l.weight.data.half() 21 | if l.bias is not None: 22 | l.bias.data = l.bias.data.half() 23 | 24 | 25 | def convert_module_to_f32(l): 26 | """ 27 | Convert primitive modules to float32, undoing convert_module_to_f16(). 28 | """ 29 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 30 | l.weight.data = l.weight.data.float() 31 | if l.bias is not None: 32 | l.bias.data = l.bias.data.float() 33 | 34 | 35 | def make_master_params(param_groups_and_shapes): 36 | """ 37 | Copy model parameters into a (differently-shaped) list of full-precision 38 | parameters. 39 | """ 40 | master_params = [] 41 | for param_group, shape in param_groups_and_shapes: 42 | master_param = nn.Parameter( 43 | _flatten_dense_tensors( 44 | [param.detach().float() for (_, param) in param_group] 45 | ).view(shape) 46 | ) 47 | master_param.requires_grad = True 48 | master_params.append(master_param) 49 | return master_params 50 | 51 | 52 | def model_grads_to_master_grads(param_groups_and_shapes, master_params): 53 | """ 54 | Copy the gradients from the model parameters into the master parameters 55 | from make_master_params(). 56 | """ 57 | for master_param, (param_group, shape) in zip( 58 | master_params, param_groups_and_shapes 59 | ): 60 | master_param.grad = _flatten_dense_tensors( 61 | [param_grad_or_zeros(param) for (_, param) in param_group] 62 | ).view(shape) 63 | 64 | 65 | def master_params_to_model_params(param_groups_and_shapes, master_params): 66 | """ 67 | Copy the master parameter data back into the model parameters. 68 | """ 69 | # Without copying to a list, if a generator is passed, this will 70 | # silently not copy any parameters. 71 | for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes): 72 | for (_, param), unflat_master_param in zip( 73 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 74 | ): 75 | param.detach().copy_(unflat_master_param) 76 | 77 | 78 | def unflatten_master_params(param_group, master_param): 79 | return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group]) 80 | 81 | 82 | def get_param_groups_and_shapes(named_model_params): 83 | named_model_params = list(named_model_params) 84 | scalar_vector_named_params = ( 85 | [(n, p) for (n, p) in named_model_params if p.ndim <= 1], 86 | (-1), 87 | ) 88 | matrix_named_params = ( 89 | [(n, p) for (n, p) in named_model_params if p.ndim > 1], 90 | (1, -1), 91 | ) 92 | return [scalar_vector_named_params, matrix_named_params] 93 | 94 | 95 | def master_params_to_state_dict( 96 | model, param_groups_and_shapes, master_params, use_fp16 97 | ): 98 | if use_fp16: 99 | state_dict = model.state_dict() 100 | for master_param, (param_group, _) in zip( 101 | master_params, param_groups_and_shapes 102 | ): 103 | for (name, _), unflat_master_param in zip( 104 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 105 | ): 106 | assert name in state_dict 107 | state_dict[name] = unflat_master_param 108 | else: 109 | state_dict = model.state_dict() 110 | for i, (name, _value) in enumerate(model.named_parameters()): 111 | assert name in state_dict 112 | state_dict[name] = master_params[i] 113 | return state_dict 114 | 115 | 116 | def state_dict_to_master_params(model, state_dict, use_fp16): 117 | if use_fp16: 118 | named_model_params = [ 119 | (name, state_dict[name]) for name, _ in model.named_parameters() 120 | ] 121 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 122 | master_params = make_master_params(param_groups_and_shapes) 123 | else: 124 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 125 | return master_params 126 | 127 | 128 | def zero_master_grads(master_params): 129 | for param in master_params: 130 | param.grad = None 131 | 132 | 133 | def zero_grad(model_params): 134 | for param in model_params: 135 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 136 | if param.grad is not None: 137 | param.grad.detach_() 138 | param.grad.zero_() 139 | 140 | 141 | def param_grad_or_zeros(param): 142 | if param.grad is not None: 143 | return param.grad.data.detach() 144 | else: 145 | return th.zeros_like(param) 146 | 147 | 148 | class MixedPrecisionTrainer: 149 | def __init__( 150 | self, 151 | *, 152 | model, 153 | use_fp16=False, 154 | fp16_scale_growth=1e-3, 155 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 156 | ): 157 | self.model = model 158 | self.use_fp16 = use_fp16 159 | self.fp16_scale_growth = fp16_scale_growth 160 | 161 | self.model_params = list(self.model.parameters()) 162 | self.master_params = self.model_params 163 | self.param_groups_and_shapes = None 164 | self.lg_loss_scale = initial_lg_loss_scale 165 | 166 | if self.use_fp16: 167 | self.param_groups_and_shapes = get_param_groups_and_shapes( 168 | self.model.named_parameters() 169 | ) 170 | self.master_params = make_master_params(self.param_groups_and_shapes) 171 | self.model.convert_to_fp16() 172 | 173 | def zero_grad(self): 174 | zero_grad(self.model_params) 175 | 176 | def backward(self, loss: th.Tensor): 177 | if self.use_fp16: 178 | loss_scale = 2 ** self.lg_loss_scale 179 | (loss * loss_scale).backward() 180 | else: 181 | loss.backward() 182 | 183 | def optimize(self, opt: th.optim.Optimizer): 184 | if self.use_fp16: 185 | return self._optimize_fp16(opt) 186 | else: 187 | return self._optimize_normal(opt) 188 | 189 | def _optimize_fp16(self, opt: th.optim.Optimizer): 190 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 191 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 192 | grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale) 193 | if check_overflow(grad_norm): 194 | self.lg_loss_scale -= 1 195 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 196 | zero_master_grads(self.master_params) 197 | return False 198 | 199 | logger.logkv_mean("grad_norm", grad_norm) 200 | logger.logkv_mean("param_norm", param_norm) 201 | 202 | self.master_params[0].grad.mul_(1.0 / (2 ** self.lg_loss_scale)) 203 | opt.step() 204 | zero_master_grads(self.master_params) 205 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 206 | self.lg_loss_scale += self.fp16_scale_growth 207 | return True 208 | 209 | def _optimize_normal(self, opt: th.optim.Optimizer): 210 | grad_norm, param_norm = self._compute_norms() 211 | logger.logkv_mean("grad_norm", grad_norm) 212 | logger.logkv_mean("param_norm", param_norm) 213 | opt.step() 214 | return True 215 | 216 | def _compute_norms(self, grad_scale=1.0): 217 | grad_norm = 0.0 218 | param_norm = 0.0 219 | for p in self.master_params: 220 | with th.no_grad(): 221 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 222 | if p.grad is not None: 223 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 224 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 225 | 226 | def master_params_to_state_dict(self, master_params): 227 | return master_params_to_state_dict( 228 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 229 | ) 230 | 231 | def state_dict_to_master_params(self, state_dict): 232 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 233 | 234 | 235 | def check_overflow(value): 236 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 237 | -------------------------------------------------------------------------------- /guided_diffusion/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import torch as th 7 | import torch.distributed as dist 8 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 9 | from torch.optim import AdamW 10 | 11 | from . import dist_util, logger 12 | from .fp16_util import MixedPrecisionTrainer 13 | from .nn import update_ema 14 | from .resample import LossAwareSampler, UniformSampler 15 | 16 | # For ImageNet experiments, this was a good default value. 17 | # We found that the lg_loss_scale quickly climbed to 18 | # 20-21 within the first ~1K steps of training. 19 | INITIAL_LOG_LOSS_SCALE = 20.0 20 | 21 | 22 | class TrainLoop: 23 | def __init__( 24 | self, 25 | *, 26 | model, 27 | diffusion, 28 | data, 29 | batch_size, 30 | microbatch, 31 | lr, 32 | ema_rate, 33 | log_interval, 34 | save_interval, 35 | resume_checkpoint, 36 | use_fp16=False, 37 | fp16_scale_growth=1e-3, 38 | schedule_sampler=None, 39 | weight_decay=0.0, 40 | lr_anneal_steps=0, 41 | ): 42 | self.model = model 43 | self.diffusion = diffusion 44 | self.data = data 45 | self.batch_size = batch_size 46 | self.microbatch = microbatch if microbatch > 0 else batch_size 47 | self.lr = lr 48 | self.ema_rate = ( 49 | [ema_rate] 50 | if isinstance(ema_rate, float) 51 | else [float(x) for x in ema_rate.split(",")] 52 | ) 53 | self.log_interval = log_interval 54 | self.save_interval = save_interval 55 | self.resume_checkpoint = resume_checkpoint 56 | self.use_fp16 = use_fp16 57 | self.fp16_scale_growth = fp16_scale_growth 58 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 59 | self.weight_decay = weight_decay 60 | self.lr_anneal_steps = lr_anneal_steps 61 | 62 | self.step = 0 63 | self.resume_step = 0 64 | self.global_batch = self.batch_size * dist.get_world_size() 65 | 66 | self.sync_cuda = th.cuda.is_available() 67 | 68 | self._load_and_sync_parameters() 69 | self.mp_trainer = MixedPrecisionTrainer( 70 | model=self.model, 71 | use_fp16=self.use_fp16, 72 | fp16_scale_growth=fp16_scale_growth, 73 | ) 74 | 75 | self.opt = AdamW( 76 | self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay 77 | ) 78 | if self.resume_step: 79 | self._load_optimizer_state() 80 | # Model was resumed, either due to a restart or a checkpoint 81 | # being specified at the command line. 82 | self.ema_params = [ 83 | self._load_ema_parameters(rate) for rate in self.ema_rate 84 | ] 85 | else: 86 | self.ema_params = [ 87 | copy.deepcopy(self.mp_trainer.master_params) 88 | for _ in range(len(self.ema_rate)) 89 | ] 90 | 91 | if th.cuda.is_available(): 92 | self.use_ddp = True 93 | self.ddp_model = DDP( 94 | self.model, 95 | device_ids=[dist_util.dev()], 96 | output_device=dist_util.dev(), 97 | broadcast_buffers=False, 98 | bucket_cap_mb=128, 99 | find_unused_parameters=False, 100 | ) 101 | else: 102 | if dist.get_world_size() > 1: 103 | logger.warn( 104 | "Distributed training requires CUDA. " 105 | "Gradients will not be synchronized properly!" 106 | ) 107 | self.use_ddp = False 108 | self.ddp_model = self.model 109 | 110 | def _load_and_sync_parameters(self): 111 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 112 | 113 | if resume_checkpoint: 114 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 115 | if dist.get_rank() == 0: 116 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 117 | self.model.load_state_dict( 118 | dist_util.load_state_dict( 119 | resume_checkpoint, map_location=dist_util.dev() 120 | ) 121 | ) 122 | 123 | dist_util.sync_params(self.model.parameters()) 124 | 125 | def _load_ema_parameters(self, rate): 126 | ema_params = copy.deepcopy(self.mp_trainer.master_params) 127 | 128 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 129 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 130 | if ema_checkpoint: 131 | if dist.get_rank() == 0: 132 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 133 | state_dict = dist_util.load_state_dict( 134 | ema_checkpoint, map_location=dist_util.dev() 135 | ) 136 | ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) 137 | 138 | dist_util.sync_params(ema_params) 139 | return ema_params 140 | 141 | def _load_optimizer_state(self): 142 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 143 | opt_checkpoint = bf.join( 144 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 145 | ) 146 | if bf.exists(opt_checkpoint): 147 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 148 | state_dict = dist_util.load_state_dict( 149 | opt_checkpoint, map_location=dist_util.dev() 150 | ) 151 | self.opt.load_state_dict(state_dict) 152 | 153 | def run_loop(self): 154 | while ( 155 | not self.lr_anneal_steps 156 | or self.step + self.resume_step < self.lr_anneal_steps 157 | ): 158 | batch, cond = next(self.data) 159 | self.run_step(batch, cond) 160 | if self.step % self.log_interval == 0: 161 | logger.dumpkvs() 162 | if self.step % self.save_interval == 0: 163 | self.save() 164 | # Run for a finite amount of time in integration tests. 165 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 166 | return 167 | self.step += 1 168 | # Save the last checkpoint if it wasn't already saved. 169 | if (self.step - 1) % self.save_interval != 0: 170 | self.save() 171 | 172 | def run_step(self, batch, cond): 173 | self.forward_backward(batch, cond) 174 | took_step = self.mp_trainer.optimize(self.opt) 175 | if took_step: 176 | self._update_ema() 177 | self._anneal_lr() 178 | self.log_step() 179 | 180 | def forward_backward(self, batch, cond): 181 | self.mp_trainer.zero_grad() 182 | for i in range(0, batch.shape[0], self.microbatch): 183 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 184 | micro_cond = { 185 | k: v[i : i + self.microbatch].to(dist_util.dev()) 186 | for k, v in cond.items() 187 | } 188 | last_batch = (i + self.microbatch) >= batch.shape[0] 189 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 190 | 191 | compute_losses = functools.partial( 192 | self.diffusion.training_losses, 193 | self.ddp_model, 194 | micro, 195 | t, 196 | model_kwargs=micro_cond, 197 | ) 198 | 199 | if last_batch or not self.use_ddp: 200 | losses = compute_losses() 201 | else: 202 | with self.ddp_model.no_sync(): 203 | losses = compute_losses() 204 | 205 | if isinstance(self.schedule_sampler, LossAwareSampler): 206 | self.schedule_sampler.update_with_local_losses( 207 | t, losses["loss"].detach() 208 | ) 209 | 210 | loss = (losses["loss"] * weights).mean() 211 | log_loss_dict( 212 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 213 | ) 214 | self.mp_trainer.backward(loss) 215 | 216 | def _update_ema(self): 217 | for rate, params in zip(self.ema_rate, self.ema_params): 218 | update_ema(params, self.mp_trainer.master_params, rate=rate) 219 | 220 | def _anneal_lr(self): 221 | if not self.lr_anneal_steps: 222 | return 223 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 224 | lr = self.lr * (1 - frac_done) 225 | for param_group in self.opt.param_groups: 226 | param_group["lr"] = lr 227 | 228 | def log_step(self): 229 | logger.logkv("step", self.step + self.resume_step) 230 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 231 | 232 | def save(self): 233 | def save_checkpoint(rate, params): 234 | state_dict = self.mp_trainer.master_params_to_state_dict(params) 235 | if dist.get_rank() == 0: 236 | logger.log(f"saving model {rate}...") 237 | if not rate: 238 | filename = f"model{(self.step+self.resume_step):06d}.pt" 239 | else: 240 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 241 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 242 | th.save(state_dict, f) 243 | 244 | save_checkpoint(0, self.mp_trainer.master_params) 245 | for rate, params in zip(self.ema_rate, self.ema_params): 246 | save_checkpoint(rate, params) 247 | 248 | if dist.get_rank() == 0: 249 | with bf.BlobFile( 250 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 251 | "wb", 252 | ) as f: 253 | th.save(self.opt.state_dict(), f) 254 | 255 | dist.barrier() 256 | 257 | 258 | def parse_resume_step_from_filename(filename): 259 | """ 260 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 261 | checkpoint's number of steps. 262 | """ 263 | split = filename.split("model") 264 | if len(split) < 2: 265 | return 0 266 | split1 = split[-1].split(".")[0] 267 | try: 268 | return int(split1) 269 | except ValueError: 270 | return 0 271 | 272 | 273 | def get_blob_logdir(): 274 | # You can change this to be a separate path to save checkpoints to 275 | # a blobstore or some external drive. 276 | return logger.get_dir() 277 | 278 | 279 | def find_resume_checkpoint(): 280 | # On your infrastructure, you may want to override this to automatically 281 | # discover the latest checkpoint on your blob storage, etc. 282 | return None 283 | 284 | 285 | def find_ema_checkpoint(main_checkpoint, step, rate): 286 | if main_checkpoint is None: 287 | return None 288 | filename = f"ema_{rate}_{(step):06d}.pt" 289 | path = bf.join(bf.dirname(main_checkpoint), filename) 290 | if bf.exists(path): 291 | return path 292 | return None 293 | 294 | 295 | def log_loss_dict(diffusion, ts, losses): 296 | for key, values in losses.items(): 297 | logger.logkv_mean(key, values.mean().item()) 298 | # Log the quantiles (four quartiles, in particular). 299 | for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 300 | quartile = int(4 * sub_t / diffusion.num_timesteps) 301 | logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 302 | -------------------------------------------------------------------------------- /guided_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, EncoderUNetModel 7 | 8 | NUM_CLASSES = 1000 9 | 10 | 11 | def diffusion_defaults(): 12 | """ 13 | Defaults for image and classifier training. 14 | """ 15 | return dict( 16 | learn_sigma=False, 17 | diffusion_steps=1000, 18 | noise_schedule="linear", 19 | timestep_respacing="", 20 | use_kl=False, 21 | predict_xstart=False, 22 | rescale_timesteps=False, 23 | rescale_learned_sigmas=False, 24 | ) 25 | 26 | 27 | def classifier_defaults(): 28 | """ 29 | Defaults for classifier models. 30 | """ 31 | return dict( 32 | image_size=64, 33 | classifier_use_fp16=False, 34 | classifier_width=128, 35 | classifier_depth=2, 36 | classifier_attention_resolutions="32,16,8", # 16 37 | classifier_use_scale_shift_norm=True, # False 38 | classifier_resblock_updown=True, # False 39 | classifier_pool="attention", 40 | ) 41 | 42 | 43 | def model_and_diffusion_defaults(): 44 | """ 45 | Defaults for image training. 46 | """ 47 | res = dict( 48 | image_size=64, 49 | num_channels=128, 50 | num_res_blocks=2, 51 | num_heads=4, 52 | num_heads_upsample=-1, 53 | num_head_channels=-1, 54 | attention_resolutions="16,8", 55 | channel_mult="", 56 | dropout=0.0, 57 | p2_gamma=0, 58 | p2_k=1, 59 | class_cond=False, 60 | use_checkpoint=False, 61 | use_scale_shift_norm=True, 62 | resblock_updown=False, 63 | use_fp16=False, 64 | use_new_attention_order=False, 65 | ) 66 | res.update(diffusion_defaults()) 67 | return res 68 | 69 | 70 | def classifier_and_diffusion_defaults(): 71 | res = classifier_defaults() 72 | res.update(diffusion_defaults()) 73 | return res 74 | 75 | 76 | def create_model_and_diffusion( 77 | image_size, 78 | class_cond, 79 | learn_sigma, 80 | num_channels, 81 | num_res_blocks, 82 | channel_mult, 83 | num_heads, 84 | num_head_channels, 85 | num_heads_upsample, 86 | attention_resolutions, 87 | dropout, 88 | p2_gamma, 89 | p2_k, 90 | diffusion_steps, 91 | noise_schedule, 92 | timestep_respacing, 93 | use_kl, 94 | predict_xstart, 95 | rescale_timesteps, 96 | rescale_learned_sigmas, 97 | use_checkpoint, 98 | use_scale_shift_norm, 99 | resblock_updown, 100 | use_fp16, 101 | use_new_attention_order, 102 | ): 103 | model = create_model( 104 | image_size, 105 | num_channels, 106 | num_res_blocks, 107 | channel_mult=channel_mult, 108 | learn_sigma=learn_sigma, 109 | class_cond=class_cond, 110 | use_checkpoint=use_checkpoint, 111 | attention_resolutions=attention_resolutions, 112 | num_heads=num_heads, 113 | num_head_channels=num_head_channels, 114 | num_heads_upsample=num_heads_upsample, 115 | use_scale_shift_norm=use_scale_shift_norm, 116 | dropout=dropout, 117 | resblock_updown=resblock_updown, 118 | use_fp16=use_fp16, 119 | use_new_attention_order=use_new_attention_order, 120 | ) 121 | diffusion = create_gaussian_diffusion( 122 | steps=diffusion_steps, 123 | learn_sigma=learn_sigma, 124 | noise_schedule=noise_schedule, 125 | use_kl=use_kl, 126 | predict_xstart=predict_xstart, 127 | rescale_timesteps=rescale_timesteps, 128 | rescale_learned_sigmas=rescale_learned_sigmas, 129 | timestep_respacing=timestep_respacing, 130 | p2_gamma=p2_gamma, 131 | p2_k=p2_k, 132 | ) 133 | return model, diffusion 134 | 135 | 136 | def create_model( 137 | image_size, 138 | num_channels, 139 | num_res_blocks, 140 | channel_mult="", 141 | learn_sigma=False, 142 | class_cond=False, 143 | use_checkpoint=False, 144 | attention_resolutions="16", 145 | num_heads=1, 146 | num_head_channels=-1, 147 | num_heads_upsample=-1, 148 | use_scale_shift_norm=False, 149 | dropout=0, 150 | resblock_updown=False, 151 | use_fp16=False, 152 | use_new_attention_order=False, 153 | ): 154 | if channel_mult == "": 155 | if image_size == 512: 156 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 157 | elif image_size == 256: 158 | channel_mult = (1, 1, 2, 2, 4, 4) 159 | elif image_size == 128: 160 | channel_mult = (1, 1, 2, 3, 4) 161 | elif image_size == 64: 162 | channel_mult = (1, 2, 3, 4) 163 | else: 164 | raise ValueError(f"unsupported image size: {image_size}") 165 | else: 166 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 167 | 168 | attention_ds = [] 169 | for res in attention_resolutions.split(","): 170 | attention_ds.append(image_size // int(res)) 171 | 172 | return UNetModel( 173 | image_size=image_size, 174 | in_channels=3, 175 | model_channels=num_channels, 176 | out_channels=(3 if not learn_sigma else 6), 177 | num_res_blocks=num_res_blocks, 178 | attention_resolutions=tuple(attention_ds), 179 | dropout=dropout, 180 | channel_mult=channel_mult, 181 | num_classes=(NUM_CLASSES if class_cond else None), 182 | use_checkpoint=use_checkpoint, 183 | use_fp16=use_fp16, 184 | num_heads=num_heads, 185 | num_head_channels=num_head_channels, 186 | num_heads_upsample=num_heads_upsample, 187 | use_scale_shift_norm=use_scale_shift_norm, 188 | resblock_updown=resblock_updown, 189 | use_new_attention_order=use_new_attention_order, 190 | ) 191 | 192 | 193 | def create_classifier_and_diffusion( 194 | image_size, 195 | classifier_use_fp16, 196 | classifier_width, 197 | classifier_depth, 198 | classifier_attention_resolutions, 199 | classifier_use_scale_shift_norm, 200 | classifier_resblock_updown, 201 | classifier_pool, 202 | learn_sigma, 203 | diffusion_steps, 204 | noise_schedule, 205 | timestep_respacing, 206 | use_kl, 207 | predict_xstart, 208 | rescale_timesteps, 209 | rescale_learned_sigmas, 210 | ): 211 | classifier = create_classifier( 212 | image_size, 213 | classifier_use_fp16, 214 | classifier_width, 215 | classifier_depth, 216 | classifier_attention_resolutions, 217 | classifier_use_scale_shift_norm, 218 | classifier_resblock_updown, 219 | classifier_pool, 220 | ) 221 | diffusion = create_gaussian_diffusion( 222 | steps=diffusion_steps, 223 | learn_sigma=learn_sigma, 224 | noise_schedule=noise_schedule, 225 | use_kl=use_kl, 226 | predict_xstart=predict_xstart, 227 | rescale_timesteps=rescale_timesteps, 228 | rescale_learned_sigmas=rescale_learned_sigmas, 229 | timestep_respacing=timestep_respacing, 230 | ) 231 | return classifier, diffusion 232 | 233 | 234 | def create_classifier( 235 | image_size, 236 | classifier_use_fp16, 237 | classifier_width, 238 | classifier_depth, 239 | classifier_attention_resolutions, 240 | classifier_use_scale_shift_norm, 241 | classifier_resblock_updown, 242 | classifier_pool, 243 | ): 244 | if image_size == 512: 245 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 246 | elif image_size == 256: 247 | channel_mult = (1, 1, 2, 2, 4, 4) 248 | elif image_size == 128: 249 | channel_mult = (1, 1, 2, 3, 4) 250 | elif image_size == 64: 251 | channel_mult = (1, 2, 3, 4) 252 | else: 253 | raise ValueError(f"unsupported image size: {image_size}") 254 | 255 | attention_ds = [] 256 | for res in classifier_attention_resolutions.split(","): 257 | attention_ds.append(image_size // int(res)) 258 | 259 | return EncoderUNetModel( 260 | image_size=image_size, 261 | in_channels=3, 262 | model_channels=classifier_width, 263 | out_channels=1000, 264 | num_res_blocks=classifier_depth, 265 | attention_resolutions=tuple(attention_ds), 266 | channel_mult=channel_mult, 267 | use_fp16=classifier_use_fp16, 268 | num_head_channels=64, 269 | use_scale_shift_norm=classifier_use_scale_shift_norm, 270 | resblock_updown=classifier_resblock_updown, 271 | pool=classifier_pool, 272 | ) 273 | 274 | 275 | def sr_model_and_diffusion_defaults(): 276 | res = model_and_diffusion_defaults() 277 | res["large_size"] = 256 278 | res["small_size"] = 64 279 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 280 | for k in res.copy().keys(): 281 | if k not in arg_names: 282 | del res[k] 283 | return res 284 | 285 | 286 | def sr_create_model_and_diffusion( 287 | large_size, 288 | small_size, 289 | class_cond, 290 | learn_sigma, 291 | num_channels, 292 | num_res_blocks, 293 | num_heads, 294 | num_head_channels, 295 | num_heads_upsample, 296 | attention_resolutions, 297 | dropout, 298 | diffusion_steps, 299 | noise_schedule, 300 | timestep_respacing, 301 | use_kl, 302 | predict_xstart, 303 | rescale_timesteps, 304 | rescale_learned_sigmas, 305 | use_checkpoint, 306 | use_scale_shift_norm, 307 | resblock_updown, 308 | use_fp16, 309 | p2_gamma, 310 | p2_k, 311 | ): 312 | model = sr_create_model( 313 | large_size, 314 | small_size, 315 | num_channels, 316 | num_res_blocks, 317 | learn_sigma=learn_sigma, 318 | class_cond=class_cond, 319 | use_checkpoint=use_checkpoint, 320 | attention_resolutions=attention_resolutions, 321 | num_heads=num_heads, 322 | num_head_channels=num_head_channels, 323 | num_heads_upsample=num_heads_upsample, 324 | use_scale_shift_norm=use_scale_shift_norm, 325 | dropout=dropout, 326 | resblock_updown=resblock_updown, 327 | use_fp16=use_fp16, 328 | ) 329 | diffusion = create_gaussian_diffusion( 330 | steps=diffusion_steps, 331 | learn_sigma=learn_sigma, 332 | noise_schedule=noise_schedule, 333 | use_kl=use_kl, 334 | predict_xstart=predict_xstart, 335 | rescale_timesteps=rescale_timesteps, 336 | rescale_learned_sigmas=rescale_learned_sigmas, 337 | timestep_respacing=timestep_respacing, 338 | p2_gamma=p2_gamma, 339 | p2_k=p2_k, 340 | ) 341 | return model, diffusion 342 | 343 | 344 | def sr_create_model( 345 | large_size, 346 | small_size, 347 | num_channels, 348 | num_res_blocks, 349 | learn_sigma, 350 | class_cond, 351 | use_checkpoint, 352 | attention_resolutions, 353 | num_heads, 354 | num_head_channels, 355 | num_heads_upsample, 356 | use_scale_shift_norm, 357 | dropout, 358 | resblock_updown, 359 | use_fp16, 360 | ): 361 | _ = small_size # hack to prevent unused variable 362 | 363 | if large_size == 512: 364 | channel_mult = (1, 1, 2, 2, 4, 4) 365 | elif large_size == 256: 366 | channel_mult = (1, 1, 2, 2, 4, 4) 367 | elif large_size == 64: 368 | channel_mult = (1, 2, 3, 4) 369 | else: 370 | raise ValueError(f"unsupported large size: {large_size}") 371 | 372 | attention_ds = [] 373 | for res in attention_resolutions.split(","): 374 | attention_ds.append(large_size // int(res)) 375 | 376 | return SuperResModel( 377 | image_size=large_size, 378 | in_channels=3, 379 | model_channels=num_channels, 380 | out_channels=(3 if not learn_sigma else 6), 381 | num_res_blocks=num_res_blocks, 382 | attention_resolutions=tuple(attention_ds), 383 | dropout=dropout, 384 | channel_mult=channel_mult, 385 | num_classes=(NUM_CLASSES if class_cond else None), 386 | use_checkpoint=use_checkpoint, 387 | num_heads=num_heads, 388 | num_head_channels=num_head_channels, 389 | num_heads_upsample=num_heads_upsample, 390 | use_scale_shift_norm=use_scale_shift_norm, 391 | resblock_updown=resblock_updown, 392 | use_fp16=use_fp16, 393 | ) 394 | 395 | 396 | def create_gaussian_diffusion( 397 | *, 398 | steps=1000, 399 | learn_sigma=False, 400 | sigma_small=False, 401 | noise_schedule="linear", 402 | use_kl=False, 403 | predict_xstart=False, 404 | rescale_timesteps=False, 405 | rescale_learned_sigmas=False, 406 | timestep_respacing="", 407 | p2_gamma=0, 408 | p2_k=1, 409 | ): 410 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 411 | if use_kl: 412 | loss_type = gd.LossType.RESCALED_KL 413 | elif rescale_learned_sigmas: 414 | loss_type = gd.LossType.RESCALED_MSE 415 | else: 416 | loss_type = gd.LossType.MSE 417 | if not timestep_respacing: 418 | timestep_respacing = [steps] 419 | return SpacedDiffusion( 420 | use_timesteps=space_timesteps(steps, timestep_respacing), 421 | betas=betas, 422 | model_mean_type=( 423 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 424 | ), 425 | model_var_type=( 426 | ( 427 | gd.ModelVarType.FIXED_LARGE 428 | if not sigma_small 429 | else gd.ModelVarType.FIXED_SMALL 430 | ) 431 | if not learn_sigma 432 | else gd.ModelVarType.LEARNED_RANGE 433 | ), 434 | loss_type=loss_type, 435 | rescale_timesteps=rescale_timesteps, 436 | p2_gamma=p2_gamma, 437 | p2_k=p2_k, 438 | ) 439 | 440 | 441 | def add_dict_to_argparser(parser, default_dict): 442 | for k, v in default_dict.items(): 443 | v_type = type(v) 444 | if v is None: 445 | v_type = str 446 | elif isinstance(v, bool): 447 | v_type = str2bool 448 | parser.add_argument(f"--{k}", default=v, type=v_type) 449 | 450 | 451 | def args_to_dict(args, keys): 452 | return {k: getattr(args, k) for k in keys} 453 | 454 | 455 | def str2bool(v): 456 | """ 457 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 458 | """ 459 | if isinstance(v, bool): 460 | return v 461 | if v.lower() in ("yes", "true", "t", "y", "1"): 462 | return True 463 | elif v.lower() in ("no", "false", "f", "n", "0"): 464 | return False 465 | else: 466 | raise argparse.ArgumentTypeError("boolean value expected") 467 | -------------------------------------------------------------------------------- /guided_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 | -------------------------------------------------------------------------------- /evaluations/evaluator.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import io 3 | import os 4 | import random 5 | import warnings 6 | import zipfile 7 | from abc import ABC, abstractmethod 8 | from contextlib import contextmanager 9 | from functools import partial 10 | from multiprocessing import cpu_count 11 | from multiprocessing.pool import ThreadPool 12 | from typing import Iterable, Optional, Tuple 13 | 14 | import numpy as np 15 | import requests 16 | import tensorflow.compat.v1 as tf 17 | from scipy import linalg 18 | from tqdm.auto import tqdm 19 | 20 | INCEPTION_V3_URL = "https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/classify_image_graph_def.pb" 21 | INCEPTION_V3_PATH = "classify_image_graph_def.pb" 22 | 23 | FID_POOL_NAME = "pool_3:0" 24 | FID_SPATIAL_NAME = "mixed_6/conv:0" 25 | 26 | 27 | def main(): 28 | parser = argparse.ArgumentParser() 29 | parser.add_argument("ref_batch", help="path to reference batch npz file") 30 | parser.add_argument("sample_batch", help="path to sample batch npz file") 31 | args = parser.parse_args() 32 | 33 | config = tf.ConfigProto( 34 | allow_soft_placement=True # allows DecodeJpeg to run on CPU in Inception graph 35 | ) 36 | config.gpu_options.allow_growth = True 37 | evaluator = Evaluator(tf.Session(config=config)) 38 | 39 | print("warming up TensorFlow...") 40 | # This will cause TF to print a bunch of verbose stuff now rather 41 | # than after the next print(), to help prevent confusion. 42 | evaluator.warmup() 43 | 44 | print("computing reference batch activations...") 45 | ref_acts = evaluator.read_activations(args.ref_batch) 46 | print("computing/reading reference batch statistics...") 47 | ref_stats, ref_stats_spatial = evaluator.read_statistics(args.ref_batch, ref_acts) 48 | 49 | print("computing sample batch activations...") 50 | sample_acts = evaluator.read_activations(args.sample_batch) 51 | print("computing/reading sample batch statistics...") 52 | sample_stats, sample_stats_spatial = evaluator.read_statistics(args.sample_batch, sample_acts) 53 | 54 | print("Computing evaluations...") 55 | print("Inception Score:", evaluator.compute_inception_score(sample_acts[0])) 56 | print("FID:", sample_stats.frechet_distance(ref_stats)) 57 | print("sFID:", sample_stats_spatial.frechet_distance(ref_stats_spatial)) 58 | prec, recall = evaluator.compute_prec_recall(ref_acts[0], sample_acts[0]) 59 | print("Precision:", prec) 60 | print("Recall:", recall) 61 | 62 | 63 | class InvalidFIDException(Exception): 64 | pass 65 | 66 | 67 | class FIDStatistics: 68 | def __init__(self, mu: np.ndarray, sigma: np.ndarray): 69 | self.mu = mu 70 | self.sigma = sigma 71 | 72 | def frechet_distance(self, other, eps=1e-6): 73 | """ 74 | Compute the Frechet distance between two sets of statistics. 75 | """ 76 | # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L132 77 | mu1, sigma1 = self.mu, self.sigma 78 | mu2, sigma2 = other.mu, other.sigma 79 | 80 | mu1 = np.atleast_1d(mu1) 81 | mu2 = np.atleast_1d(mu2) 82 | 83 | sigma1 = np.atleast_2d(sigma1) 84 | sigma2 = np.atleast_2d(sigma2) 85 | 86 | assert ( 87 | mu1.shape == mu2.shape 88 | ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}" 89 | assert ( 90 | sigma1.shape == sigma2.shape 91 | ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}" 92 | 93 | diff = mu1 - mu2 94 | 95 | # product might be almost singular 96 | covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) 97 | if not np.isfinite(covmean).all(): 98 | msg = ( 99 | "fid calculation produces singular product; adding %s to diagonal of cov estimates" 100 | % eps 101 | ) 102 | warnings.warn(msg) 103 | offset = np.eye(sigma1.shape[0]) * eps 104 | covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) 105 | 106 | # numerical error might give slight imaginary component 107 | if np.iscomplexobj(covmean): 108 | if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): 109 | m = np.max(np.abs(covmean.imag)) 110 | raise ValueError("Imaginary component {}".format(m)) 111 | covmean = covmean.real 112 | 113 | tr_covmean = np.trace(covmean) 114 | 115 | return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean 116 | 117 | 118 | class Evaluator: 119 | def __init__( 120 | self, 121 | session, 122 | batch_size=64, 123 | softmax_batch_size=512, 124 | ): 125 | self.sess = session 126 | self.batch_size = batch_size 127 | self.softmax_batch_size = softmax_batch_size 128 | self.manifold_estimator = ManifoldEstimator(session) 129 | with self.sess.graph.as_default(): 130 | self.image_input = tf.placeholder(tf.float32, shape=[None, None, None, 3]) 131 | self.softmax_input = tf.placeholder(tf.float32, shape=[None, 2048]) 132 | self.pool_features, self.spatial_features = _create_feature_graph(self.image_input) 133 | self.softmax = _create_softmax_graph(self.softmax_input) 134 | 135 | def warmup(self): 136 | self.compute_activations(np.zeros([1, 8, 64, 64, 3])) 137 | 138 | def read_activations(self, npz_path: str) -> Tuple[np.ndarray, np.ndarray]: 139 | with open_npz_array(npz_path, "arr_0") as reader: 140 | return self.compute_activations(reader.read_batches(self.batch_size)) 141 | 142 | def compute_activations(self, batches: Iterable[np.ndarray]) -> Tuple[np.ndarray, np.ndarray]: 143 | """ 144 | Compute image features for downstream evals. 145 | 146 | :param batches: a iterator over NHWC numpy arrays in [0, 255]. 147 | :return: a tuple of numpy arrays of shape [N x X], where X is a feature 148 | dimension. The tuple is (pool_3, spatial). 149 | """ 150 | preds = [] 151 | spatial_preds = [] 152 | for batch in tqdm(batches): 153 | batch = batch.astype(np.float32) 154 | pred, spatial_pred = self.sess.run( 155 | [self.pool_features, self.spatial_features], {self.image_input: batch} 156 | ) 157 | preds.append(pred.reshape([pred.shape[0], -1])) 158 | spatial_preds.append(spatial_pred.reshape([spatial_pred.shape[0], -1])) 159 | return ( 160 | np.concatenate(preds, axis=0), 161 | np.concatenate(spatial_preds, axis=0), 162 | ) 163 | 164 | def read_statistics( 165 | self, npz_path: str, activations: Tuple[np.ndarray, np.ndarray] 166 | ) -> Tuple[FIDStatistics, FIDStatistics]: 167 | obj = np.load(npz_path) 168 | if "mu" in list(obj.keys()): 169 | return FIDStatistics(obj["mu"], obj["sigma"]), FIDStatistics( 170 | obj["mu_s"], obj["sigma_s"] 171 | ) 172 | return tuple(self.compute_statistics(x) for x in activations) 173 | 174 | def compute_statistics(self, activations: np.ndarray) -> FIDStatistics: 175 | mu = np.mean(activations, axis=0) 176 | sigma = np.cov(activations, rowvar=False) 177 | return FIDStatistics(mu, sigma) 178 | 179 | def compute_inception_score(self, activations: np.ndarray, split_size: int = 5000) -> float: 180 | softmax_out = [] 181 | for i in range(0, len(activations), self.softmax_batch_size): 182 | acts = activations[i : i + self.softmax_batch_size] 183 | softmax_out.append(self.sess.run(self.softmax, feed_dict={self.softmax_input: acts})) 184 | preds = np.concatenate(softmax_out, axis=0) 185 | # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46 186 | scores = [] 187 | for i in range(0, len(preds), split_size): 188 | part = preds[i : i + split_size] 189 | kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))) 190 | kl = np.mean(np.sum(kl, 1)) 191 | scores.append(np.exp(kl)) 192 | return float(np.mean(scores)) 193 | 194 | def compute_prec_recall( 195 | self, activations_ref: np.ndarray, activations_sample: np.ndarray 196 | ) -> Tuple[float, float]: 197 | radii_1 = self.manifold_estimator.manifold_radii(activations_ref) 198 | radii_2 = self.manifold_estimator.manifold_radii(activations_sample) 199 | pr = self.manifold_estimator.evaluate_pr( 200 | activations_ref, radii_1, activations_sample, radii_2 201 | ) 202 | return (float(pr[0][0]), float(pr[1][0])) 203 | 204 | 205 | class ManifoldEstimator: 206 | """ 207 | A helper for comparing manifolds of feature vectors. 208 | 209 | Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L57 210 | """ 211 | 212 | def __init__( 213 | self, 214 | session, 215 | row_batch_size=10000, 216 | col_batch_size=10000, 217 | nhood_sizes=(3,), 218 | clamp_to_percentile=None, 219 | eps=1e-5, 220 | ): 221 | """ 222 | Estimate the manifold of given feature vectors. 223 | 224 | :param session: the TensorFlow session. 225 | :param row_batch_size: row batch size to compute pairwise distances 226 | (parameter to trade-off between memory usage and performance). 227 | :param col_batch_size: column batch size to compute pairwise distances. 228 | :param nhood_sizes: number of neighbors used to estimate the manifold. 229 | :param clamp_to_percentile: prune hyperspheres that have radius larger than 230 | the given percentile. 231 | :param eps: small number for numerical stability. 232 | """ 233 | self.distance_block = DistanceBlock(session) 234 | self.row_batch_size = row_batch_size 235 | self.col_batch_size = col_batch_size 236 | self.nhood_sizes = nhood_sizes 237 | self.num_nhoods = len(nhood_sizes) 238 | self.clamp_to_percentile = clamp_to_percentile 239 | self.eps = eps 240 | 241 | def warmup(self): 242 | feats, radii = ( 243 | np.zeros([1, 2048], dtype=np.float32), 244 | np.zeros([1, 1], dtype=np.float32), 245 | ) 246 | self.evaluate_pr(feats, radii, feats, radii) 247 | 248 | def manifold_radii(self, features: np.ndarray) -> np.ndarray: 249 | num_images = len(features) 250 | 251 | # Estimate manifold of features by calculating distances to k-NN of each sample. 252 | radii = np.zeros([num_images, self.num_nhoods], dtype=np.float32) 253 | distance_batch = np.zeros([self.row_batch_size, num_images], dtype=np.float32) 254 | seq = np.arange(max(self.nhood_sizes) + 1, dtype=np.int32) 255 | 256 | for begin1 in range(0, num_images, self.row_batch_size): 257 | end1 = min(begin1 + self.row_batch_size, num_images) 258 | row_batch = features[begin1:end1] 259 | 260 | for begin2 in range(0, num_images, self.col_batch_size): 261 | end2 = min(begin2 + self.col_batch_size, num_images) 262 | col_batch = features[begin2:end2] 263 | 264 | # Compute distances between batches. 265 | distance_batch[ 266 | 0 : end1 - begin1, begin2:end2 267 | ] = self.distance_block.pairwise_distances(row_batch, col_batch) 268 | 269 | # Find the k-nearest neighbor from the current batch. 270 | radii[begin1:end1, :] = np.concatenate( 271 | [ 272 | x[:, self.nhood_sizes] 273 | for x in _numpy_partition(distance_batch[0 : end1 - begin1, :], seq, axis=1) 274 | ], 275 | axis=0, 276 | ) 277 | 278 | if self.clamp_to_percentile is not None: 279 | max_distances = np.percentile(radii, self.clamp_to_percentile, axis=0) 280 | radii[radii > max_distances] = 0 281 | return radii 282 | 283 | def evaluate(self, features: np.ndarray, radii: np.ndarray, eval_features: np.ndarray): 284 | """ 285 | Evaluate if new feature vectors are at the manifold. 286 | """ 287 | num_eval_images = eval_features.shape[0] 288 | num_ref_images = radii.shape[0] 289 | distance_batch = np.zeros([self.row_batch_size, num_ref_images], dtype=np.float32) 290 | batch_predictions = np.zeros([num_eval_images, self.num_nhoods], dtype=np.int32) 291 | max_realism_score = np.zeros([num_eval_images], dtype=np.float32) 292 | nearest_indices = np.zeros([num_eval_images], dtype=np.int32) 293 | 294 | for begin1 in range(0, num_eval_images, self.row_batch_size): 295 | end1 = min(begin1 + self.row_batch_size, num_eval_images) 296 | feature_batch = eval_features[begin1:end1] 297 | 298 | for begin2 in range(0, num_ref_images, self.col_batch_size): 299 | end2 = min(begin2 + self.col_batch_size, num_ref_images) 300 | ref_batch = features[begin2:end2] 301 | 302 | distance_batch[ 303 | 0 : end1 - begin1, begin2:end2 304 | ] = self.distance_block.pairwise_distances(feature_batch, ref_batch) 305 | 306 | # From the minibatch of new feature vectors, determine if they are in the estimated manifold. 307 | # If a feature vector is inside a hypersphere of some reference sample, then 308 | # the new sample lies at the estimated manifold. 309 | # The radii of the hyperspheres are determined from distances of neighborhood size k. 310 | samples_in_manifold = distance_batch[0 : end1 - begin1, :, None] <= radii 311 | batch_predictions[begin1:end1] = np.any(samples_in_manifold, axis=1).astype(np.int32) 312 | 313 | max_realism_score[begin1:end1] = np.max( 314 | radii[:, 0] / (distance_batch[0 : end1 - begin1, :] + self.eps), axis=1 315 | ) 316 | nearest_indices[begin1:end1] = np.argmin(distance_batch[0 : end1 - begin1, :], axis=1) 317 | 318 | return { 319 | "fraction": float(np.mean(batch_predictions)), 320 | "batch_predictions": batch_predictions, 321 | "max_realisim_score": max_realism_score, 322 | "nearest_indices": nearest_indices, 323 | } 324 | 325 | def evaluate_pr( 326 | self, 327 | features_1: np.ndarray, 328 | radii_1: np.ndarray, 329 | features_2: np.ndarray, 330 | radii_2: np.ndarray, 331 | ) -> Tuple[np.ndarray, np.ndarray]: 332 | """ 333 | Evaluate precision and recall efficiently. 334 | 335 | :param features_1: [N1 x D] feature vectors for reference batch. 336 | :param radii_1: [N1 x K1] radii for reference vectors. 337 | :param features_2: [N2 x D] feature vectors for the other batch. 338 | :param radii_2: [N x K2] radii for other vectors. 339 | :return: a tuple of arrays for (precision, recall): 340 | - precision: an np.ndarray of length K1 341 | - recall: an np.ndarray of length K2 342 | """ 343 | features_1_status = np.zeros([len(features_1), radii_2.shape[1]], dtype=np.bool) 344 | features_2_status = np.zeros([len(features_2), radii_1.shape[1]], dtype=np.bool) 345 | for begin_1 in range(0, len(features_1), self.row_batch_size): 346 | end_1 = begin_1 + self.row_batch_size 347 | batch_1 = features_1[begin_1:end_1] 348 | for begin_2 in range(0, len(features_2), self.col_batch_size): 349 | end_2 = begin_2 + self.col_batch_size 350 | batch_2 = features_2[begin_2:end_2] 351 | batch_1_in, batch_2_in = self.distance_block.less_thans( 352 | batch_1, radii_1[begin_1:end_1], batch_2, radii_2[begin_2:end_2] 353 | ) 354 | features_1_status[begin_1:end_1] |= batch_1_in 355 | features_2_status[begin_2:end_2] |= batch_2_in 356 | return ( 357 | np.mean(features_2_status.astype(np.float64), axis=0), 358 | np.mean(features_1_status.astype(np.float64), axis=0), 359 | ) 360 | 361 | 362 | class DistanceBlock: 363 | """ 364 | Calculate pairwise distances between vectors. 365 | 366 | Adapted from https://github.com/kynkaat/improved-precision-and-recall-metric/blob/f60f25e5ad933a79135c783fcda53de30f42c9b9/precision_recall.py#L34 367 | """ 368 | 369 | def __init__(self, session): 370 | self.session = session 371 | 372 | # Initialize TF graph to calculate pairwise distances. 373 | with session.graph.as_default(): 374 | self._features_batch1 = tf.placeholder(tf.float32, shape=[None, None]) 375 | self._features_batch2 = tf.placeholder(tf.float32, shape=[None, None]) 376 | distance_block_16 = _batch_pairwise_distances( 377 | tf.cast(self._features_batch1, tf.float16), 378 | tf.cast(self._features_batch2, tf.float16), 379 | ) 380 | self.distance_block = tf.cond( 381 | tf.reduce_all(tf.math.is_finite(distance_block_16)), 382 | lambda: tf.cast(distance_block_16, tf.float32), 383 | lambda: _batch_pairwise_distances(self._features_batch1, self._features_batch2), 384 | ) 385 | 386 | # Extra logic for less thans. 387 | self._radii1 = tf.placeholder(tf.float32, shape=[None, None]) 388 | self._radii2 = tf.placeholder(tf.float32, shape=[None, None]) 389 | dist32 = tf.cast(self.distance_block, tf.float32)[..., None] 390 | self._batch_1_in = tf.math.reduce_any(dist32 <= self._radii2, axis=1) 391 | self._batch_2_in = tf.math.reduce_any(dist32 <= self._radii1[:, None], axis=0) 392 | 393 | def pairwise_distances(self, U, V): 394 | """ 395 | Evaluate pairwise distances between two batches of feature vectors. 396 | """ 397 | return self.session.run( 398 | self.distance_block, 399 | feed_dict={self._features_batch1: U, self._features_batch2: V}, 400 | ) 401 | 402 | def less_thans(self, batch_1, radii_1, batch_2, radii_2): 403 | return self.session.run( 404 | [self._batch_1_in, self._batch_2_in], 405 | feed_dict={ 406 | self._features_batch1: batch_1, 407 | self._features_batch2: batch_2, 408 | self._radii1: radii_1, 409 | self._radii2: radii_2, 410 | }, 411 | ) 412 | 413 | 414 | def _batch_pairwise_distances(U, V): 415 | """ 416 | Compute pairwise distances between two batches of feature vectors. 417 | """ 418 | with tf.variable_scope("pairwise_dist_block"): 419 | # Squared norms of each row in U and V. 420 | norm_u = tf.reduce_sum(tf.square(U), 1) 421 | norm_v = tf.reduce_sum(tf.square(V), 1) 422 | 423 | # norm_u as a column and norm_v as a row vectors. 424 | norm_u = tf.reshape(norm_u, [-1, 1]) 425 | norm_v = tf.reshape(norm_v, [1, -1]) 426 | 427 | # Pairwise squared Euclidean distances. 428 | D = tf.maximum(norm_u - 2 * tf.matmul(U, V, False, True) + norm_v, 0.0) 429 | 430 | return D 431 | 432 | 433 | class NpzArrayReader(ABC): 434 | @abstractmethod 435 | def read_batch(self, batch_size: int) -> Optional[np.ndarray]: 436 | pass 437 | 438 | @abstractmethod 439 | def remaining(self) -> int: 440 | pass 441 | 442 | def read_batches(self, batch_size: int) -> Iterable[np.ndarray]: 443 | def gen_fn(): 444 | while True: 445 | batch = self.read_batch(batch_size) 446 | if batch is None: 447 | break 448 | yield batch 449 | 450 | rem = self.remaining() 451 | num_batches = rem // batch_size + int(rem % batch_size != 0) 452 | return BatchIterator(gen_fn, num_batches) 453 | 454 | 455 | class BatchIterator: 456 | def __init__(self, gen_fn, length): 457 | self.gen_fn = gen_fn 458 | self.length = length 459 | 460 | def __len__(self): 461 | return self.length 462 | 463 | def __iter__(self): 464 | return self.gen_fn() 465 | 466 | 467 | class StreamingNpzArrayReader(NpzArrayReader): 468 | def __init__(self, arr_f, shape, dtype): 469 | self.arr_f = arr_f 470 | self.shape = shape 471 | self.dtype = dtype 472 | self.idx = 0 473 | 474 | def read_batch(self, batch_size: int) -> Optional[np.ndarray]: 475 | if self.idx >= self.shape[0]: 476 | return None 477 | 478 | bs = min(batch_size, self.shape[0] - self.idx) 479 | self.idx += bs 480 | 481 | if self.dtype.itemsize == 0: 482 | return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype) 483 | 484 | read_count = bs * np.prod(self.shape[1:]) 485 | read_size = int(read_count * self.dtype.itemsize) 486 | data = _read_bytes(self.arr_f, read_size, "array data") 487 | return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]]) 488 | 489 | def remaining(self) -> int: 490 | return max(0, self.shape[0] - self.idx) 491 | 492 | 493 | class MemoryNpzArrayReader(NpzArrayReader): 494 | def __init__(self, arr): 495 | self.arr = arr 496 | self.idx = 0 497 | 498 | @classmethod 499 | def load(cls, path: str, arr_name: str): 500 | with open(path, "rb") as f: 501 | arr = np.load(f)[arr_name] 502 | return cls(arr) 503 | 504 | def read_batch(self, batch_size: int) -> Optional[np.ndarray]: 505 | if self.idx >= self.arr.shape[0]: 506 | return None 507 | 508 | res = self.arr[self.idx : self.idx + batch_size] 509 | self.idx += batch_size 510 | return res 511 | 512 | def remaining(self) -> int: 513 | return max(0, self.arr.shape[0] - self.idx) 514 | 515 | 516 | @contextmanager 517 | def open_npz_array(path: str, arr_name: str) -> NpzArrayReader: 518 | with _open_npy_file(path, arr_name) as arr_f: 519 | version = np.lib.format.read_magic(arr_f) 520 | if version == (1, 0): 521 | header = np.lib.format.read_array_header_1_0(arr_f) 522 | elif version == (2, 0): 523 | header = np.lib.format.read_array_header_2_0(arr_f) 524 | else: 525 | yield MemoryNpzArrayReader.load(path, arr_name) 526 | return 527 | shape, fortran, dtype = header 528 | if fortran or dtype.hasobject: 529 | yield MemoryNpzArrayReader.load(path, arr_name) 530 | else: 531 | yield StreamingNpzArrayReader(arr_f, shape, dtype) 532 | 533 | 534 | def _read_bytes(fp, size, error_template="ran out of data"): 535 | """ 536 | Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886 537 | 538 | Read from file-like object until size bytes are read. 539 | Raises ValueError if not EOF is encountered before size bytes are read. 540 | Non-blocking objects only supported if they derive from io objects. 541 | Required as e.g. ZipExtFile in python 2.6 can return less data than 542 | requested. 543 | """ 544 | data = bytes() 545 | while True: 546 | # io files (default in python3) return None or raise on 547 | # would-block, python2 file will truncate, probably nothing can be 548 | # done about that. note that regular files can't be non-blocking 549 | try: 550 | r = fp.read(size - len(data)) 551 | data += r 552 | if len(r) == 0 or len(data) == size: 553 | break 554 | except io.BlockingIOError: 555 | pass 556 | if len(data) != size: 557 | msg = "EOF: reading %s, expected %d bytes got %d" 558 | raise ValueError(msg % (error_template, size, len(data))) 559 | else: 560 | return data 561 | 562 | 563 | @contextmanager 564 | def _open_npy_file(path: str, arr_name: str): 565 | with open(path, "rb") as f: 566 | with zipfile.ZipFile(f, "r") as zip_f: 567 | if f"{arr_name}.npy" not in zip_f.namelist(): 568 | raise ValueError(f"missing {arr_name} in npz file") 569 | with zip_f.open(f"{arr_name}.npy", "r") as arr_f: 570 | yield arr_f 571 | 572 | 573 | def _download_inception_model(): 574 | if os.path.exists(INCEPTION_V3_PATH): 575 | return 576 | print("downloading InceptionV3 model...") 577 | with requests.get(INCEPTION_V3_URL, stream=True) as r: 578 | r.raise_for_status() 579 | tmp_path = INCEPTION_V3_PATH + ".tmp" 580 | with open(tmp_path, "wb") as f: 581 | for chunk in tqdm(r.iter_content(chunk_size=8192)): 582 | f.write(chunk) 583 | os.rename(tmp_path, INCEPTION_V3_PATH) 584 | 585 | 586 | def _create_feature_graph(input_batch): 587 | _download_inception_model() 588 | prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}" 589 | with open(INCEPTION_V3_PATH, "rb") as f: 590 | graph_def = tf.GraphDef() 591 | graph_def.ParseFromString(f.read()) 592 | pool3, spatial = tf.import_graph_def( 593 | graph_def, 594 | input_map={f"ExpandDims:0": input_batch}, 595 | return_elements=[FID_POOL_NAME, FID_SPATIAL_NAME], 596 | name=prefix, 597 | ) 598 | _update_shapes(pool3) 599 | spatial = spatial[..., :7] 600 | return pool3, spatial 601 | 602 | 603 | def _create_softmax_graph(input_batch): 604 | _download_inception_model() 605 | prefix = f"{random.randrange(2**32)}_{random.randrange(2**32)}" 606 | with open(INCEPTION_V3_PATH, "rb") as f: 607 | graph_def = tf.GraphDef() 608 | graph_def.ParseFromString(f.read()) 609 | (matmul,) = tf.import_graph_def( 610 | graph_def, return_elements=[f"softmax/logits/MatMul"], name=prefix 611 | ) 612 | w = matmul.inputs[1] 613 | logits = tf.matmul(input_batch, w) 614 | return tf.nn.softmax(logits) 615 | 616 | 617 | def _update_shapes(pool3): 618 | # https://github.com/bioinf-jku/TTUR/blob/73ab375cdf952a12686d9aa7978567771084da42/fid.py#L50-L63 619 | ops = pool3.graph.get_operations() 620 | for op in ops: 621 | for o in op.outputs: 622 | shape = o.get_shape() 623 | if shape._dims is not None: # pylint: disable=protected-access 624 | # shape = [s.value for s in shape] TF 1.x 625 | shape = [s for s in shape] # TF 2.x 626 | new_shape = [] 627 | for j, s in enumerate(shape): 628 | if s == 1 and j == 0: 629 | new_shape.append(None) 630 | else: 631 | new_shape.append(s) 632 | o.__dict__["_shape_val"] = tf.TensorShape(new_shape) 633 | return pool3 634 | 635 | 636 | def _numpy_partition(arr, kth, **kwargs): 637 | num_workers = min(cpu_count(), len(arr)) 638 | chunk_size = len(arr) // num_workers 639 | extra = len(arr) % num_workers 640 | 641 | start_idx = 0 642 | batches = [] 643 | for i in range(num_workers): 644 | size = chunk_size + (1 if i < extra else 0) 645 | batches.append(arr[start_idx : start_idx + size]) 646 | start_idx += size 647 | 648 | with ThreadPool(num_workers) as pool: 649 | return list(pool.map(partial(np.partition, kth=kth, **kwargs), batches)) 650 | 651 | 652 | if __name__ == "__main__": 653 | main() 654 | -------------------------------------------------------------------------------- /guided_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 | checkpoint, 13 | conv_nd, 14 | linear, 15 | avg_pool_nd, 16 | zero_module, 17 | normalization, 18 | timestep_embedding, 19 | ) 20 | 21 | 22 | class AttentionPool2d(nn.Module): 23 | """ 24 | Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py 25 | """ 26 | 27 | def __init__( 28 | self, 29 | spacial_dim: int, 30 | embed_dim: int, 31 | num_heads_channels: int, 32 | output_dim: int = None, 33 | ): 34 | super().__init__() 35 | self.positional_embedding = nn.Parameter( 36 | th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5 37 | ) 38 | self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1) 39 | self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1) 40 | self.num_heads = embed_dim // num_heads_channels 41 | self.attention = QKVAttention(self.num_heads) 42 | 43 | def forward(self, x): 44 | b, c, *_spatial = x.shape 45 | x = x.reshape(b, c, -1) # NC(HW) 46 | x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1) 47 | x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1) 48 | x = self.qkv_proj(x) 49 | x = self.attention(x) 50 | x = self.c_proj(x) 51 | return x[:, :, 0] 52 | 53 | 54 | class TimestepBlock(nn.Module): 55 | """ 56 | Any module where forward() takes timestep embeddings as a second argument. 57 | """ 58 | 59 | @abstractmethod 60 | def forward(self, x, emb): 61 | """ 62 | Apply the module to `x` given `emb` timestep embeddings. 63 | """ 64 | 65 | 66 | class TimestepEmbedSequential(nn.Sequential, TimestepBlock): 67 | """ 68 | A sequential module that passes timestep embeddings to the children that 69 | support it as an extra input. 70 | """ 71 | 72 | def forward(self, x, emb): 73 | for layer in self: 74 | if isinstance(layer, TimestepBlock): 75 | x = layer(x, emb) 76 | else: 77 | x = layer(x) 78 | return x 79 | 80 | 81 | class Upsample(nn.Module): 82 | """ 83 | An upsampling 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 | upsampling occurs in the inner-two dimensions. 89 | """ 90 | 91 | def __init__(self, channels, use_conv, dims=2, out_channels=None): 92 | super().__init__() 93 | self.channels = channels 94 | self.out_channels = out_channels or channels 95 | self.use_conv = use_conv 96 | self.dims = dims 97 | if use_conv: 98 | self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1) 99 | 100 | def forward(self, x): 101 | assert x.shape[1] == self.channels 102 | if self.dims == 3: 103 | x = F.interpolate( 104 | x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" 105 | ) 106 | else: 107 | x = F.interpolate(x, scale_factor=2, mode="nearest") 108 | if self.use_conv: 109 | x = self.conv(x) 110 | return x 111 | 112 | 113 | class Downsample(nn.Module): 114 | """ 115 | A downsampling layer with an optional convolution. 116 | 117 | :param channels: channels in the inputs and outputs. 118 | :param use_conv: a bool determining if a convolution is applied. 119 | :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then 120 | downsampling occurs in the inner-two dimensions. 121 | """ 122 | 123 | def __init__(self, channels, use_conv, dims=2, out_channels=None): 124 | super().__init__() 125 | self.channels = channels 126 | self.out_channels = out_channels or channels 127 | self.use_conv = use_conv 128 | self.dims = dims 129 | stride = 2 if dims != 3 else (1, 2, 2) 130 | if use_conv: 131 | self.op = conv_nd( 132 | dims, self.channels, self.out_channels, 3, stride=stride, padding=1 133 | ) 134 | else: 135 | assert self.channels == self.out_channels 136 | self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) 137 | 138 | def forward(self, x): 139 | assert x.shape[1] == self.channels 140 | return self.op(x) 141 | 142 | 143 | class ResBlock(TimestepBlock): 144 | """ 145 | A residual block that can optionally change the number of channels. 146 | 147 | :param channels: the number of input channels. 148 | :param emb_channels: the number of timestep embedding channels. 149 | :param dropout: the rate of dropout. 150 | :param out_channels: if specified, the number of out channels. 151 | :param use_conv: if True and out_channels is specified, use a spatial 152 | convolution instead of a smaller 1x1 convolution to change the 153 | channels in the skip connection. 154 | :param dims: determines if the signal is 1D, 2D, or 3D. 155 | :param use_checkpoint: if True, use gradient checkpointing on this module. 156 | :param up: if True, use this block for upsampling. 157 | :param down: if True, use this block for downsampling. 158 | """ 159 | 160 | def __init__( 161 | self, 162 | channels, 163 | emb_channels, 164 | dropout, 165 | out_channels=None, 166 | use_conv=False, 167 | use_scale_shift_norm=False, 168 | dims=2, 169 | use_checkpoint=False, 170 | up=False, 171 | down=False, 172 | ): 173 | super().__init__() 174 | self.channels = channels 175 | self.emb_channels = emb_channels 176 | self.dropout = dropout 177 | self.out_channels = out_channels or channels 178 | self.use_conv = use_conv 179 | self.use_checkpoint = use_checkpoint 180 | self.use_scale_shift_norm = use_scale_shift_norm 181 | 182 | self.in_layers = nn.Sequential( 183 | normalization(channels), 184 | nn.SiLU(), 185 | conv_nd(dims, channels, self.out_channels, 3, padding=1), 186 | ) 187 | 188 | self.updown = up or down 189 | 190 | if up: 191 | self.h_upd = Upsample(channels, False, dims) 192 | self.x_upd = Upsample(channels, False, dims) 193 | elif down: 194 | self.h_upd = Downsample(channels, False, dims) 195 | self.x_upd = Downsample(channels, False, dims) 196 | else: 197 | self.h_upd = self.x_upd = nn.Identity() 198 | 199 | self.emb_layers = nn.Sequential( 200 | nn.SiLU(), 201 | linear( 202 | emb_channels, 203 | 2 * self.out_channels if use_scale_shift_norm else self.out_channels, 204 | ), 205 | ) 206 | self.out_layers = nn.Sequential( 207 | normalization(self.out_channels), 208 | nn.SiLU(), 209 | nn.Dropout(p=dropout), 210 | zero_module( 211 | conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) 212 | ), 213 | ) 214 | 215 | if self.out_channels == channels: 216 | self.skip_connection = nn.Identity() 217 | elif use_conv: 218 | self.skip_connection = conv_nd( 219 | dims, channels, self.out_channels, 3, padding=1 220 | ) 221 | else: 222 | self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) 223 | 224 | def forward(self, x, emb): 225 | """ 226 | Apply the block to a Tensor, conditioned on a timestep embedding. 227 | 228 | :param x: an [N x C x ...] Tensor of features. 229 | :param emb: an [N x emb_channels] Tensor of timestep embeddings. 230 | :return: an [N x C x ...] Tensor of outputs. 231 | """ 232 | return checkpoint( 233 | self._forward, (x, emb), self.parameters(), self.use_checkpoint 234 | ) 235 | 236 | def _forward(self, x, emb): 237 | if self.updown: 238 | in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] 239 | h = in_rest(x) 240 | h = self.h_upd(h) 241 | x = self.x_upd(x) 242 | h = in_conv(h) 243 | else: 244 | h = self.in_layers(x) 245 | emb_out = self.emb_layers(emb).type(h.dtype) 246 | while len(emb_out.shape) < len(h.shape): 247 | emb_out = emb_out[..., None] 248 | if self.use_scale_shift_norm: 249 | out_norm, out_rest = self.out_layers[0], self.out_layers[1:] 250 | scale, shift = th.chunk(emb_out, 2, dim=1) 251 | h = out_norm(h) * (1 + scale) + shift 252 | h = out_rest(h) 253 | else: 254 | h = h + emb_out 255 | h = self.out_layers(h) 256 | return self.skip_connection(x) + h 257 | 258 | 259 | class AttentionBlock(nn.Module): 260 | """ 261 | An attention block that allows spatial positions to attend to each other. 262 | 263 | Originally ported from here, but adapted to the N-d case. 264 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. 265 | """ 266 | 267 | def __init__( 268 | self, 269 | channels, 270 | num_heads=1, 271 | num_head_channels=-1, 272 | use_checkpoint=False, 273 | use_new_attention_order=False, 274 | ): 275 | super().__init__() 276 | self.channels = channels 277 | if num_head_channels == -1: 278 | self.num_heads = num_heads 279 | else: 280 | assert ( 281 | channels % num_head_channels == 0 282 | ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" 283 | self.num_heads = channels // num_head_channels 284 | self.use_checkpoint = use_checkpoint 285 | self.norm = normalization(channels) 286 | self.qkv = conv_nd(1, channels, channels * 3, 1) 287 | if use_new_attention_order: 288 | # split qkv before split heads 289 | self.attention = QKVAttention(self.num_heads) 290 | else: 291 | # split heads before split qkv 292 | self.attention = QKVAttentionLegacy(self.num_heads) 293 | 294 | self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) 295 | 296 | def forward(self, x): 297 | return checkpoint(self._forward, (x,), self.parameters(), True) 298 | 299 | def _forward(self, x): 300 | b, c, *spatial = x.shape 301 | x = x.reshape(b, c, -1) 302 | qkv = self.qkv(self.norm(x)) 303 | h = self.attention(qkv) 304 | h = self.proj_out(h) 305 | return (x + h).reshape(b, c, *spatial) 306 | 307 | 308 | def count_flops_attn(model, _x, y): 309 | """ 310 | A counter for the `thop` package to count the operations in an 311 | attention operation. 312 | Meant to be used like: 313 | macs, params = thop.profile( 314 | model, 315 | inputs=(inputs, timestamps), 316 | custom_ops={QKVAttention: QKVAttention.count_flops}, 317 | ) 318 | """ 319 | b, c, *spatial = y[0].shape 320 | num_spatial = int(np.prod(spatial)) 321 | # We perform two matmuls with the same number of ops. 322 | # The first computes the weight matrix, the second computes 323 | # the combination of the value vectors. 324 | matmul_ops = 2 * b * (num_spatial ** 2) * c 325 | model.total_ops += th.DoubleTensor([matmul_ops]) 326 | 327 | 328 | class QKVAttentionLegacy(nn.Module): 329 | """ 330 | A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping 331 | """ 332 | 333 | def __init__(self, n_heads): 334 | super().__init__() 335 | self.n_heads = n_heads 336 | 337 | def forward(self, qkv): 338 | """ 339 | Apply QKV attention. 340 | 341 | :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. 342 | :return: an [N x (H * C) x T] tensor after attention. 343 | """ 344 | bs, width, length = qkv.shape 345 | assert width % (3 * self.n_heads) == 0 346 | ch = width // (3 * self.n_heads) 347 | q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1) 348 | scale = 1 / math.sqrt(math.sqrt(ch)) 349 | weight = th.einsum( 350 | "bct,bcs->bts", q * scale, k * scale 351 | ) # More stable with f16 than dividing afterwards 352 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) 353 | a = th.einsum("bts,bcs->bct", weight, v) 354 | return a.reshape(bs, -1, length) 355 | 356 | @staticmethod 357 | def count_flops(model, _x, y): 358 | return count_flops_attn(model, _x, y) 359 | 360 | 361 | class QKVAttention(nn.Module): 362 | """ 363 | A module which performs QKV attention and splits in a different order. 364 | """ 365 | 366 | def __init__(self, n_heads): 367 | super().__init__() 368 | self.n_heads = n_heads 369 | 370 | def forward(self, qkv): 371 | """ 372 | Apply QKV attention. 373 | 374 | :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. 375 | :return: an [N x (H * C) x T] tensor after attention. 376 | """ 377 | bs, width, length = qkv.shape 378 | assert width % (3 * self.n_heads) == 0 379 | ch = width // (3 * self.n_heads) 380 | q, k, v = qkv.chunk(3, dim=1) 381 | scale = 1 / math.sqrt(math.sqrt(ch)) 382 | weight = th.einsum( 383 | "bct,bcs->bts", 384 | (q * scale).view(bs * self.n_heads, ch, length), 385 | (k * scale).view(bs * self.n_heads, ch, length), 386 | ) # More stable with f16 than dividing afterwards 387 | weight = th.softmax(weight.float(), dim=-1).type(weight.dtype) 388 | a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length)) 389 | return a.reshape(bs, -1, length) 390 | 391 | @staticmethod 392 | def count_flops(model, _x, y): 393 | return count_flops_attn(model, _x, y) 394 | 395 | 396 | class UNetModel(nn.Module): 397 | """ 398 | The full UNet model with attention and timestep embedding. 399 | 400 | :param in_channels: channels in the input Tensor. 401 | :param model_channels: base channel count for the model. 402 | :param out_channels: channels in the output Tensor. 403 | :param num_res_blocks: number of residual blocks per downsample. 404 | :param attention_resolutions: a collection of downsample rates at which 405 | attention will take place. May be a set, list, or tuple. 406 | For example, if this contains 4, then at 4x downsampling, attention 407 | will be used. 408 | :param dropout: the dropout probability. 409 | :param channel_mult: channel multiplier for each level of the UNet. 410 | :param conv_resample: if True, use learned convolutions for upsampling and 411 | downsampling. 412 | :param dims: determines if the signal is 1D, 2D, or 3D. 413 | :param num_classes: if specified (as an int), then this model will be 414 | class-conditional with `num_classes` classes. 415 | :param use_checkpoint: use gradient checkpointing to reduce memory usage. 416 | :param num_heads: the number of attention heads in each attention layer. 417 | :param num_heads_channels: if specified, ignore num_heads and instead use 418 | a fixed channel width per attention head. 419 | :param num_heads_upsample: works with num_heads to set a different number 420 | of heads for upsampling. Deprecated. 421 | :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. 422 | :param resblock_updown: use residual blocks for up/downsampling. 423 | :param use_new_attention_order: use a different attention pattern for potentially 424 | increased efficiency. 425 | """ 426 | 427 | def __init__( 428 | self, 429 | image_size, 430 | in_channels, 431 | model_channels, 432 | out_channels, 433 | num_res_blocks, 434 | attention_resolutions, 435 | dropout=0, 436 | channel_mult=(1, 2, 4, 8), 437 | conv_resample=True, 438 | dims=2, 439 | num_classes=None, 440 | use_checkpoint=False, 441 | use_fp16=False, 442 | num_heads=1, 443 | num_head_channels=-1, 444 | num_heads_upsample=-1, 445 | use_scale_shift_norm=False, 446 | resblock_updown=False, 447 | use_new_attention_order=False, 448 | ): 449 | super().__init__() 450 | 451 | if num_heads_upsample == -1: 452 | num_heads_upsample = num_heads 453 | 454 | self.image_size = image_size 455 | self.in_channels = in_channels 456 | self.model_channels = model_channels 457 | self.out_channels = out_channels 458 | self.num_res_blocks = num_res_blocks 459 | self.attention_resolutions = attention_resolutions 460 | self.dropout = dropout 461 | self.channel_mult = channel_mult 462 | self.conv_resample = conv_resample 463 | self.num_classes = num_classes 464 | self.use_checkpoint = use_checkpoint 465 | self.dtype = th.float16 if use_fp16 else th.float32 466 | self.num_heads = num_heads 467 | self.num_head_channels = num_head_channels 468 | self.num_heads_upsample = num_heads_upsample 469 | 470 | time_embed_dim = model_channels * 4 471 | self.time_embed = nn.Sequential( 472 | linear(model_channels, time_embed_dim), 473 | nn.SiLU(), 474 | linear(time_embed_dim, time_embed_dim), 475 | ) 476 | 477 | if self.num_classes is not None: 478 | self.label_emb = nn.Embedding(num_classes, time_embed_dim) 479 | 480 | ch = input_ch = int(channel_mult[0] * model_channels) 481 | self.input_blocks = nn.ModuleList( 482 | [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] 483 | ) 484 | self._feature_size = ch 485 | input_block_chans = [ch] 486 | ds = 1 487 | for level, mult in enumerate(channel_mult): 488 | for _ in range(num_res_blocks): 489 | layers = [ 490 | ResBlock( 491 | ch, 492 | time_embed_dim, 493 | dropout, 494 | out_channels=int(mult * model_channels), 495 | dims=dims, 496 | use_checkpoint=use_checkpoint, 497 | use_scale_shift_norm=use_scale_shift_norm, 498 | ) 499 | ] 500 | ch = int(mult * model_channels) 501 | if ds in attention_resolutions: 502 | layers.append( 503 | AttentionBlock( 504 | ch, 505 | use_checkpoint=use_checkpoint, 506 | num_heads=num_heads, 507 | num_head_channels=num_head_channels, 508 | use_new_attention_order=use_new_attention_order, 509 | ) 510 | ) 511 | self.input_blocks.append(TimestepEmbedSequential(*layers)) 512 | self._feature_size += ch 513 | input_block_chans.append(ch) 514 | if level != len(channel_mult) - 1: 515 | out_ch = ch 516 | self.input_blocks.append( 517 | TimestepEmbedSequential( 518 | ResBlock( 519 | ch, 520 | time_embed_dim, 521 | dropout, 522 | out_channels=out_ch, 523 | dims=dims, 524 | use_checkpoint=use_checkpoint, 525 | use_scale_shift_norm=use_scale_shift_norm, 526 | down=True, 527 | ) 528 | if resblock_updown 529 | else Downsample( 530 | ch, conv_resample, dims=dims, out_channels=out_ch 531 | ) 532 | ) 533 | ) 534 | ch = out_ch 535 | input_block_chans.append(ch) 536 | ds *= 2 537 | self._feature_size += ch 538 | 539 | self.middle_block = TimestepEmbedSequential( 540 | ResBlock( 541 | ch, 542 | time_embed_dim, 543 | dropout, 544 | dims=dims, 545 | use_checkpoint=use_checkpoint, 546 | use_scale_shift_norm=use_scale_shift_norm, 547 | ), 548 | AttentionBlock( 549 | ch, 550 | use_checkpoint=use_checkpoint, 551 | num_heads=num_heads, 552 | num_head_channels=num_head_channels, 553 | use_new_attention_order=use_new_attention_order, 554 | ), 555 | ResBlock( 556 | ch, 557 | time_embed_dim, 558 | dropout, 559 | dims=dims, 560 | use_checkpoint=use_checkpoint, 561 | use_scale_shift_norm=use_scale_shift_norm, 562 | ), 563 | ) 564 | self._feature_size += ch 565 | 566 | self.output_blocks = nn.ModuleList([]) 567 | for level, mult in list(enumerate(channel_mult))[::-1]: 568 | for i in range(num_res_blocks + 1): 569 | ich = input_block_chans.pop() 570 | layers = [ 571 | ResBlock( 572 | ch + ich, 573 | time_embed_dim, 574 | dropout, 575 | out_channels=int(model_channels * mult), 576 | dims=dims, 577 | use_checkpoint=use_checkpoint, 578 | use_scale_shift_norm=use_scale_shift_norm, 579 | ) 580 | ] 581 | ch = int(model_channels * mult) 582 | if ds in attention_resolutions: 583 | layers.append( 584 | AttentionBlock( 585 | ch, 586 | use_checkpoint=use_checkpoint, 587 | num_heads=num_heads_upsample, 588 | num_head_channels=num_head_channels, 589 | use_new_attention_order=use_new_attention_order, 590 | ) 591 | ) 592 | if level and i == num_res_blocks: 593 | out_ch = ch 594 | layers.append( 595 | ResBlock( 596 | ch, 597 | time_embed_dim, 598 | dropout, 599 | out_channels=out_ch, 600 | dims=dims, 601 | use_checkpoint=use_checkpoint, 602 | use_scale_shift_norm=use_scale_shift_norm, 603 | up=True, 604 | ) 605 | if resblock_updown 606 | else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) 607 | ) 608 | ds //= 2 609 | self.output_blocks.append(TimestepEmbedSequential(*layers)) 610 | self._feature_size += ch 611 | 612 | self.out = nn.Sequential( 613 | normalization(ch), 614 | nn.SiLU(), 615 | zero_module(conv_nd(dims, input_ch, out_channels, 3, padding=1)), 616 | ) 617 | 618 | def convert_to_fp16(self): 619 | """ 620 | Convert the torso of the model to float16. 621 | """ 622 | self.input_blocks.apply(convert_module_to_f16) 623 | self.middle_block.apply(convert_module_to_f16) 624 | self.output_blocks.apply(convert_module_to_f16) 625 | 626 | def convert_to_fp32(self): 627 | """ 628 | Convert the torso of the model to float32. 629 | """ 630 | self.input_blocks.apply(convert_module_to_f32) 631 | self.middle_block.apply(convert_module_to_f32) 632 | self.output_blocks.apply(convert_module_to_f32) 633 | 634 | def forward(self, x, timesteps, y=None): 635 | """ 636 | Apply the model to an input batch. 637 | 638 | :param x: an [N x C x ...] Tensor of inputs. 639 | :param timesteps: a 1-D batch of timesteps. 640 | :param y: an [N] Tensor of labels, if class-conditional. 641 | :return: an [N x C x ...] Tensor of outputs. 642 | """ 643 | assert (y is not None) == ( 644 | self.num_classes is not None 645 | ), "must specify y if and only if the model is class-conditional" 646 | 647 | hs = [] 648 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) 649 | 650 | if self.num_classes is not None: 651 | assert y.shape == (x.shape[0],) 652 | emb = emb + self.label_emb(y) 653 | 654 | h = x.type(self.dtype) 655 | for module in self.input_blocks: 656 | h = module(h, emb) 657 | hs.append(h) 658 | h = self.middle_block(h, emb) 659 | for module in self.output_blocks: 660 | h = th.cat([h, hs.pop()], dim=1) 661 | h = module(h, emb) 662 | h = h.type(x.dtype) 663 | return self.out(h) 664 | 665 | 666 | class SuperResModel(UNetModel): 667 | """ 668 | A UNetModel that performs super-resolution. 669 | 670 | Expects an extra kwarg `low_res` to condition on a low-resolution image. 671 | """ 672 | 673 | def __init__(self, image_size, in_channels, *args, **kwargs): 674 | super().__init__(image_size, in_channels * 2, *args, **kwargs) 675 | 676 | def forward(self, x, timesteps, low_res=None, **kwargs): 677 | _, _, new_height, new_width = x.shape 678 | upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear") 679 | x = th.cat([x, upsampled], dim=1) 680 | return super().forward(x, timesteps, **kwargs) 681 | 682 | 683 | class EncoderUNetModel(nn.Module): 684 | """ 685 | The half UNet model with attention and timestep embedding. 686 | 687 | For usage, see UNet. 688 | """ 689 | 690 | def __init__( 691 | self, 692 | image_size, 693 | in_channels, 694 | model_channels, 695 | out_channels, 696 | num_res_blocks, 697 | attention_resolutions, 698 | dropout=0, 699 | channel_mult=(1, 2, 4, 8), 700 | conv_resample=True, 701 | dims=2, 702 | use_checkpoint=False, 703 | use_fp16=False, 704 | num_heads=1, 705 | num_head_channels=-1, 706 | num_heads_upsample=-1, 707 | use_scale_shift_norm=False, 708 | resblock_updown=False, 709 | use_new_attention_order=False, 710 | pool="adaptive", 711 | ): 712 | super().__init__() 713 | 714 | if num_heads_upsample == -1: 715 | num_heads_upsample = num_heads 716 | 717 | self.in_channels = in_channels 718 | self.model_channels = model_channels 719 | self.out_channels = out_channels 720 | self.num_res_blocks = num_res_blocks 721 | self.attention_resolutions = attention_resolutions 722 | self.dropout = dropout 723 | self.channel_mult = channel_mult 724 | self.conv_resample = conv_resample 725 | self.use_checkpoint = use_checkpoint 726 | self.dtype = th.float16 if use_fp16 else th.float32 727 | self.num_heads = num_heads 728 | self.num_head_channels = num_head_channels 729 | self.num_heads_upsample = num_heads_upsample 730 | 731 | time_embed_dim = model_channels * 4 732 | self.time_embed = nn.Sequential( 733 | linear(model_channels, time_embed_dim), 734 | nn.SiLU(), 735 | linear(time_embed_dim, time_embed_dim), 736 | ) 737 | 738 | ch = int(channel_mult[0] * model_channels) 739 | self.input_blocks = nn.ModuleList( 740 | [TimestepEmbedSequential(conv_nd(dims, in_channels, ch, 3, padding=1))] 741 | ) 742 | self._feature_size = ch 743 | input_block_chans = [ch] 744 | ds = 1 745 | for level, mult in enumerate(channel_mult): 746 | for _ in range(num_res_blocks): 747 | layers = [ 748 | ResBlock( 749 | ch, 750 | time_embed_dim, 751 | dropout, 752 | out_channels=int(mult * model_channels), 753 | dims=dims, 754 | use_checkpoint=use_checkpoint, 755 | use_scale_shift_norm=use_scale_shift_norm, 756 | ) 757 | ] 758 | ch = int(mult * model_channels) 759 | if ds in attention_resolutions: 760 | layers.append( 761 | AttentionBlock( 762 | ch, 763 | use_checkpoint=use_checkpoint, 764 | num_heads=num_heads, 765 | num_head_channels=num_head_channels, 766 | use_new_attention_order=use_new_attention_order, 767 | ) 768 | ) 769 | self.input_blocks.append(TimestepEmbedSequential(*layers)) 770 | self._feature_size += ch 771 | input_block_chans.append(ch) 772 | if level != len(channel_mult) - 1: 773 | out_ch = ch 774 | self.input_blocks.append( 775 | TimestepEmbedSequential( 776 | ResBlock( 777 | ch, 778 | time_embed_dim, 779 | dropout, 780 | out_channels=out_ch, 781 | dims=dims, 782 | use_checkpoint=use_checkpoint, 783 | use_scale_shift_norm=use_scale_shift_norm, 784 | down=True, 785 | ) 786 | if resblock_updown 787 | else Downsample( 788 | ch, conv_resample, dims=dims, out_channels=out_ch 789 | ) 790 | ) 791 | ) 792 | ch = out_ch 793 | input_block_chans.append(ch) 794 | ds *= 2 795 | self._feature_size += ch 796 | 797 | self.middle_block = TimestepEmbedSequential( 798 | ResBlock( 799 | ch, 800 | time_embed_dim, 801 | dropout, 802 | dims=dims, 803 | use_checkpoint=use_checkpoint, 804 | use_scale_shift_norm=use_scale_shift_norm, 805 | ), 806 | AttentionBlock( 807 | ch, 808 | use_checkpoint=use_checkpoint, 809 | num_heads=num_heads, 810 | num_head_channels=num_head_channels, 811 | use_new_attention_order=use_new_attention_order, 812 | ), 813 | ResBlock( 814 | ch, 815 | time_embed_dim, 816 | dropout, 817 | dims=dims, 818 | use_checkpoint=use_checkpoint, 819 | use_scale_shift_norm=use_scale_shift_norm, 820 | ), 821 | ) 822 | self._feature_size += ch 823 | self.pool = pool 824 | if pool == "adaptive": 825 | self.out = nn.Sequential( 826 | normalization(ch), 827 | nn.SiLU(), 828 | nn.AdaptiveAvgPool2d((1, 1)), 829 | zero_module(conv_nd(dims, ch, out_channels, 1)), 830 | nn.Flatten(), 831 | ) 832 | elif pool == "attention": 833 | assert num_head_channels != -1 834 | self.out = nn.Sequential( 835 | normalization(ch), 836 | nn.SiLU(), 837 | AttentionPool2d( 838 | (image_size // ds), ch, num_head_channels, out_channels 839 | ), 840 | ) 841 | elif pool == "spatial": 842 | self.out = nn.Sequential( 843 | nn.Linear(self._feature_size, 2048), 844 | nn.ReLU(), 845 | nn.Linear(2048, self.out_channels), 846 | ) 847 | elif pool == "spatial_v2": 848 | self.out = nn.Sequential( 849 | nn.Linear(self._feature_size, 2048), 850 | normalization(2048), 851 | nn.SiLU(), 852 | nn.Linear(2048, self.out_channels), 853 | ) 854 | else: 855 | raise NotImplementedError(f"Unexpected {pool} pooling") 856 | 857 | def convert_to_fp16(self): 858 | """ 859 | Convert the torso of the model to float16. 860 | """ 861 | self.input_blocks.apply(convert_module_to_f16) 862 | self.middle_block.apply(convert_module_to_f16) 863 | 864 | def convert_to_fp32(self): 865 | """ 866 | Convert the torso of the model to float32. 867 | """ 868 | self.input_blocks.apply(convert_module_to_f32) 869 | self.middle_block.apply(convert_module_to_f32) 870 | 871 | def forward(self, x, timesteps): 872 | """ 873 | Apply the model to an input batch. 874 | 875 | :param x: an [N x C x ...] Tensor of inputs. 876 | :param timesteps: a 1-D batch of timesteps. 877 | :return: an [N x K] Tensor of outputs. 878 | """ 879 | emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) 880 | 881 | results = [] 882 | h = x.type(self.dtype) 883 | for module in self.input_blocks: 884 | h = module(h, emb) 885 | if self.pool.startswith("spatial"): 886 | results.append(h.type(x.dtype).mean(dim=(2, 3))) 887 | h = self.middle_block(h, emb) 888 | if self.pool.startswith("spatial"): 889 | results.append(h.type(x.dtype).mean(dim=(2, 3))) 890 | h = th.cat(results, axis=-1) 891 | return self.out(h) 892 | else: 893 | h = h.type(x.dtype) 894 | return self.out(h) 895 | -------------------------------------------------------------------------------- /guided_diffusion/gaussian_diffusion.py: -------------------------------------------------------------------------------- 1 | """ 2 | This code started out as a PyTorch port of Ho et al's diffusion models: 3 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py 4 | 5 | Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules. 6 | """ 7 | 8 | import enum 9 | import math 10 | 11 | import numpy as np 12 | import torch as th 13 | 14 | from .nn import mean_flat 15 | from .losses import normal_kl, discretized_gaussian_log_likelihood 16 | 17 | 18 | def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): 19 | """ 20 | Get a pre-defined beta schedule for the given name. 21 | 22 | The beta schedule library consists of beta schedules which remain similar 23 | in the limit of num_diffusion_timesteps. 24 | Beta schedules may be added, but should not be removed or changed once 25 | they are committed to maintain backwards compatibility. 26 | """ 27 | if schedule_name == "linear": 28 | # Linear schedule from Ho et al, extended to work for any number of 29 | # diffusion steps. 30 | scale = 1000 / num_diffusion_timesteps 31 | beta_start = scale * 0.0001 32 | beta_end = scale * 0.02 33 | return np.linspace( 34 | beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64 35 | ) 36 | elif schedule_name == "cosine": 37 | return betas_for_alpha_bar( 38 | num_diffusion_timesteps, 39 | lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, 40 | ) 41 | else: 42 | raise NotImplementedError(f"unknown beta schedule: {schedule_name}") 43 | 44 | 45 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): 46 | """ 47 | Create a beta schedule that discretizes the given alpha_t_bar function, 48 | which defines the cumulative product of (1-beta) over time from t = [0,1]. 49 | 50 | :param num_diffusion_timesteps: the number of betas to produce. 51 | :param alpha_bar: a lambda that takes an argument t from 0 to 1 and 52 | produces the cumulative product of (1-beta) up to that 53 | part of the diffusion process. 54 | :param max_beta: the maximum beta to use; use values lower than 1 to 55 | prevent singularities. 56 | """ 57 | betas = [] 58 | for i in range(num_diffusion_timesteps): 59 | t1 = i / num_diffusion_timesteps 60 | t2 = (i + 1) / num_diffusion_timesteps 61 | betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) 62 | return np.array(betas) 63 | 64 | 65 | class ModelMeanType(enum.Enum): 66 | """ 67 | Which type of output the model predicts. 68 | """ 69 | 70 | PREVIOUS_X = enum.auto() # the model predicts x_{t-1} 71 | START_X = enum.auto() # the model predicts x_0 72 | EPSILON = enum.auto() # the model predicts epsilon 73 | 74 | 75 | class ModelVarType(enum.Enum): 76 | """ 77 | What is used as the model's output variance. 78 | 79 | The LEARNED_RANGE option has been added to allow the model to predict 80 | values between FIXED_SMALL and FIXED_LARGE, making its job easier. 81 | """ 82 | 83 | LEARNED = enum.auto() 84 | FIXED_SMALL = enum.auto() 85 | FIXED_LARGE = enum.auto() 86 | LEARNED_RANGE = enum.auto() 87 | 88 | 89 | class LossType(enum.Enum): 90 | MSE = enum.auto() # use raw MSE loss (and KL when learning variances) 91 | RESCALED_MSE = ( 92 | enum.auto() 93 | ) # use raw MSE loss (with RESCALED_KL when learning variances) 94 | KL = enum.auto() # use the variational lower-bound 95 | RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB 96 | 97 | def is_vb(self): 98 | return self == LossType.KL or self == LossType.RESCALED_KL 99 | 100 | 101 | class GaussianDiffusion: 102 | """ 103 | Utilities for training and sampling diffusion models. 104 | 105 | Ported directly from here, and then adapted over time to further experimentation. 106 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 107 | 108 | :param betas: a 1-D numpy array of betas for each diffusion timestep, 109 | starting at T and going to 1. 110 | :param model_mean_type: a ModelMeanType determining what the model outputs. 111 | :param model_var_type: a ModelVarType determining how variance is output. 112 | :param loss_type: a LossType determining the loss function to use. 113 | :param rescale_timesteps: if True, pass floating point timesteps into the 114 | model so that they are always scaled like in the 115 | original paper (0 to 1000). 116 | """ 117 | 118 | def __init__( 119 | self, 120 | *, 121 | betas, 122 | model_mean_type, 123 | model_var_type, 124 | loss_type, 125 | rescale_timesteps=False, 126 | p2_gamma=0, 127 | p2_k=1, 128 | ): 129 | self.model_mean_type = model_mean_type 130 | self.model_var_type = model_var_type 131 | self.loss_type = loss_type 132 | self.rescale_timesteps = rescale_timesteps 133 | 134 | # Use float64 for accuracy. 135 | betas = np.array(betas, dtype=np.float64) 136 | self.betas = betas 137 | assert len(betas.shape) == 1, "betas must be 1-D" 138 | assert (betas > 0).all() and (betas <= 1).all() 139 | 140 | self.num_timesteps = int(betas.shape[0]) 141 | 142 | alphas = 1.0 - betas 143 | self.alphas_cumprod = np.cumprod(alphas, axis=0) 144 | self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) 145 | self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) 146 | assert self.alphas_cumprod_prev.shape == (self.num_timesteps,) 147 | 148 | # calculations for diffusion q(x_t | x_{t-1}) and others 149 | self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod) 150 | self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod) 151 | self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod) 152 | self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) 153 | self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1) 154 | 155 | # calculations for posterior q(x_{t-1} | x_t, x_0) 156 | self.posterior_variance = ( 157 | betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) 158 | ) 159 | # log calculation clipped because the posterior variance is 0 at the 160 | # beginning of the diffusion chain. 161 | self.posterior_log_variance_clipped = np.log( 162 | np.append(self.posterior_variance[1], self.posterior_variance[1:]) 163 | ) 164 | self.posterior_mean_coef1 = ( 165 | betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) 166 | ) 167 | self.posterior_mean_coef2 = ( 168 | (1.0 - self.alphas_cumprod_prev) 169 | * np.sqrt(alphas) 170 | / (1.0 - self.alphas_cumprod) 171 | ) 172 | 173 | # P2 weighting 174 | self.p2_gamma = p2_gamma 175 | self.p2_k = p2_k 176 | self.snr = 1.0 / (1 - self.alphas_cumprod) - 1 177 | 178 | def q_mean_variance(self, x_start, t): 179 | """ 180 | Get the distribution q(x_t | x_0). 181 | 182 | :param x_start: the [N x C x ...] tensor of noiseless inputs. 183 | :param t: the number of diffusion steps (minus 1). Here, 0 means one step. 184 | :return: A tuple (mean, variance, log_variance), all of x_start's shape. 185 | """ 186 | mean = ( 187 | _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start 188 | ) 189 | variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) 190 | log_variance = _extract_into_tensor( 191 | self.log_one_minus_alphas_cumprod, t, x_start.shape 192 | ) 193 | return mean, variance, log_variance 194 | 195 | def q_sample(self, x_start, t, noise=None): 196 | """ 197 | Diffuse the data for a given number of diffusion steps. 198 | 199 | In other words, sample from q(x_t | x_0). 200 | 201 | :param x_start: the initial data batch. 202 | :param t: the number of diffusion steps (minus 1). Here, 0 means one step. 203 | :param noise: if specified, the split-out normal noise. 204 | :return: A noisy version of x_start. 205 | """ 206 | if noise is None: 207 | noise = th.randn_like(x_start) 208 | assert noise.shape == x_start.shape 209 | return ( 210 | _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start 211 | + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) 212 | * noise 213 | ) 214 | 215 | def q_posterior_mean_variance(self, x_start, x_t, t): 216 | """ 217 | Compute the mean and variance of the diffusion posterior: 218 | 219 | q(x_{t-1} | x_t, x_0) 220 | 221 | """ 222 | assert x_start.shape == x_t.shape 223 | posterior_mean = ( 224 | _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start 225 | + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t 226 | ) 227 | posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape) 228 | posterior_log_variance_clipped = _extract_into_tensor( 229 | self.posterior_log_variance_clipped, t, x_t.shape 230 | ) 231 | assert ( 232 | posterior_mean.shape[0] 233 | == posterior_variance.shape[0] 234 | == posterior_log_variance_clipped.shape[0] 235 | == x_start.shape[0] 236 | ) 237 | return posterior_mean, posterior_variance, posterior_log_variance_clipped 238 | 239 | def p_mean_variance( 240 | self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None 241 | ): 242 | """ 243 | Apply the model to get p(x_{t-1} | x_t), as well as a prediction of 244 | the initial x, x_0. 245 | 246 | :param model: the model, which takes a signal and a batch of timesteps 247 | as input. 248 | :param x: the [N x C x ...] tensor at time t. 249 | :param t: a 1-D Tensor of timesteps. 250 | :param clip_denoised: if True, clip the denoised signal into [-1, 1]. 251 | :param denoised_fn: if not None, a function which applies to the 252 | x_start prediction before it is used to sample. Applies before 253 | clip_denoised. 254 | :param model_kwargs: if not None, a dict of extra keyword arguments to 255 | pass to the model. This can be used for conditioning. 256 | :return: a dict with the following keys: 257 | - 'mean': the model mean output. 258 | - 'variance': the model variance output. 259 | - 'log_variance': the log of 'variance'. 260 | - 'pred_xstart': the prediction for x_0. 261 | """ 262 | if model_kwargs is None: 263 | model_kwargs = {} 264 | 265 | B, C = x.shape[:2] 266 | assert t.shape == (B,) 267 | model_output = model(x, self._scale_timesteps(t), **model_kwargs) 268 | 269 | if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]: 270 | assert model_output.shape == (B, C * 2, *x.shape[2:]) 271 | model_output, model_var_values = th.split(model_output, C, dim=1) 272 | if self.model_var_type == ModelVarType.LEARNED: 273 | model_log_variance = model_var_values 274 | model_variance = th.exp(model_log_variance) 275 | else: 276 | min_log = _extract_into_tensor( 277 | self.posterior_log_variance_clipped, t, x.shape 278 | ) 279 | max_log = _extract_into_tensor(np.log(self.betas), t, x.shape) 280 | # The model_var_values is [-1, 1] for [min_var, max_var]. 281 | frac = (model_var_values + 1) / 2 282 | model_log_variance = frac * max_log + (1 - frac) * min_log 283 | model_variance = th.exp(model_log_variance) 284 | else: 285 | model_variance, model_log_variance = { 286 | # for fixedlarge, we set the initial (log-)variance like so 287 | # to get a better decoder log likelihood. 288 | ModelVarType.FIXED_LARGE: ( 289 | np.append(self.posterior_variance[1], self.betas[1:]), 290 | np.log(np.append(self.posterior_variance[1], self.betas[1:])), 291 | ), 292 | ModelVarType.FIXED_SMALL: ( 293 | self.posterior_variance, 294 | self.posterior_log_variance_clipped, 295 | ), 296 | }[self.model_var_type] 297 | model_variance = _extract_into_tensor(model_variance, t, x.shape) 298 | model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape) 299 | 300 | def process_xstart(x): 301 | if denoised_fn is not None: 302 | x = denoised_fn(x) 303 | if clip_denoised: 304 | return x.clamp(-1, 1) 305 | return x 306 | 307 | if self.model_mean_type == ModelMeanType.PREVIOUS_X: 308 | pred_xstart = process_xstart( 309 | self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output) 310 | ) 311 | model_mean = model_output 312 | elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]: 313 | if self.model_mean_type == ModelMeanType.START_X: 314 | pred_xstart = process_xstart(model_output) 315 | else: 316 | pred_xstart = process_xstart( 317 | self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output) 318 | ) 319 | model_mean, _, _ = self.q_posterior_mean_variance( 320 | x_start=pred_xstart, x_t=x, t=t 321 | ) 322 | else: 323 | raise NotImplementedError(self.model_mean_type) 324 | 325 | assert ( 326 | model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape 327 | ) 328 | return { 329 | "mean": model_mean, 330 | "variance": model_variance, 331 | "log_variance": model_log_variance, 332 | "pred_xstart": pred_xstart, 333 | } 334 | 335 | def _predict_xstart_from_eps(self, x_t, t, eps): 336 | assert x_t.shape == eps.shape 337 | return ( 338 | _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t 339 | - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps 340 | ) 341 | 342 | def _predict_xstart_from_xprev(self, x_t, t, xprev): 343 | assert x_t.shape == xprev.shape 344 | return ( # (xprev - coef2*x_t) / coef1 345 | _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev 346 | - _extract_into_tensor( 347 | self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape 348 | ) 349 | * x_t 350 | ) 351 | 352 | def _predict_eps_from_xstart(self, x_t, t, pred_xstart): 353 | return ( 354 | _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t 355 | - pred_xstart 356 | ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) 357 | 358 | def _scale_timesteps(self, t): 359 | if self.rescale_timesteps: 360 | return t.float() * (1000.0 / self.num_timesteps) 361 | return t 362 | 363 | def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): 364 | """ 365 | Compute the mean for the previous step, given a function cond_fn that 366 | computes the gradient of a conditional log probability with respect to 367 | x. In particular, cond_fn computes grad(log(p(y|x))), and we want to 368 | condition on y. 369 | 370 | This uses the conditioning strategy from Sohl-Dickstein et al. (2015). 371 | """ 372 | gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs) 373 | new_mean = ( 374 | p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() 375 | ) 376 | return new_mean 377 | 378 | def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): 379 | """ 380 | Compute what the p_mean_variance output would have been, should the 381 | model's score function be conditioned by cond_fn. 382 | 383 | See condition_mean() for details on cond_fn. 384 | 385 | Unlike condition_mean(), this instead uses the conditioning strategy 386 | from Song et al (2020). 387 | """ 388 | alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) 389 | 390 | eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"]) 391 | eps = eps - (1 - alpha_bar).sqrt() * cond_fn( 392 | x, self._scale_timesteps(t), **model_kwargs 393 | ) 394 | 395 | out = p_mean_var.copy() 396 | out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps) 397 | out["mean"], _, _ = self.q_posterior_mean_variance( 398 | x_start=out["pred_xstart"], x_t=x, t=t 399 | ) 400 | return out 401 | 402 | def p_sample( 403 | self, 404 | model, 405 | x, 406 | t, 407 | clip_denoised=True, 408 | denoised_fn=None, 409 | cond_fn=None, 410 | model_kwargs=None, 411 | ): 412 | """ 413 | Sample x_{t-1} from the model at the given timestep. 414 | 415 | :param model: the model to sample from. 416 | :param x: the current tensor at x_{t-1}. 417 | :param t: the value of t, starting at 0 for the first diffusion step. 418 | :param clip_denoised: if True, clip the x_start prediction to [-1, 1]. 419 | :param denoised_fn: if not None, a function which applies to the 420 | x_start prediction before it is used to sample. 421 | :param cond_fn: if not None, this is a gradient function that acts 422 | similarly to the model. 423 | :param model_kwargs: if not None, a dict of extra keyword arguments to 424 | pass to the model. This can be used for conditioning. 425 | :return: a dict containing the following keys: 426 | - 'sample': a random sample from the model. 427 | - 'pred_xstart': a prediction of x_0. 428 | """ 429 | out = self.p_mean_variance( 430 | model, 431 | x, 432 | t, 433 | clip_denoised=clip_denoised, 434 | denoised_fn=denoised_fn, 435 | model_kwargs=model_kwargs, 436 | ) 437 | noise = th.randn_like(x) 438 | nonzero_mask = ( 439 | (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) 440 | ) # no noise when t == 0 441 | if cond_fn is not None: 442 | out["mean"] = self.condition_mean( 443 | cond_fn, out, x, t, model_kwargs=model_kwargs 444 | ) 445 | sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise 446 | return {"sample": sample, "pred_xstart": out["pred_xstart"]} 447 | 448 | def p_sample_loop( 449 | self, 450 | model, 451 | shape, 452 | noise=None, 453 | clip_denoised=True, 454 | denoised_fn=None, 455 | cond_fn=None, 456 | model_kwargs=None, 457 | device=None, 458 | progress=False, 459 | ): 460 | """ 461 | Generate samples from the model. 462 | 463 | :param model: the model module. 464 | :param shape: the shape of the samples, (N, C, H, W). 465 | :param noise: if specified, the noise from the encoder to sample. 466 | Should be of the same shape as `shape`. 467 | :param clip_denoised: if True, clip x_start predictions to [-1, 1]. 468 | :param denoised_fn: if not None, a function which applies to the 469 | x_start prediction before it is used to sample. 470 | :param cond_fn: if not None, this is a gradient function that acts 471 | similarly to the model. 472 | :param model_kwargs: if not None, a dict of extra keyword arguments to 473 | pass to the model. This can be used for conditioning. 474 | :param device: if specified, the device to create the samples on. 475 | If not specified, use a model parameter's device. 476 | :param progress: if True, show a tqdm progress bar. 477 | :return: a non-differentiable batch of samples. 478 | """ 479 | final = None 480 | for sample in self.p_sample_loop_progressive( 481 | model, 482 | shape, 483 | noise=noise, 484 | clip_denoised=clip_denoised, 485 | denoised_fn=denoised_fn, 486 | cond_fn=cond_fn, 487 | model_kwargs=model_kwargs, 488 | device=device, 489 | progress=progress, 490 | ): 491 | final = sample 492 | return final["sample"] 493 | 494 | def p_sample_loop_progressive( 495 | self, 496 | model, 497 | shape, 498 | noise=None, 499 | clip_denoised=True, 500 | denoised_fn=None, 501 | cond_fn=None, 502 | model_kwargs=None, 503 | device=None, 504 | progress=False, 505 | ): 506 | """ 507 | Generate samples from the model and yield intermediate samples from 508 | each timestep of diffusion. 509 | 510 | Arguments are the same as p_sample_loop(). 511 | Returns a generator over dicts, where each dict is the return value of 512 | p_sample(). 513 | """ 514 | if device is None: 515 | device = next(model.parameters()).device 516 | assert isinstance(shape, (tuple, list)) 517 | if noise is not None: 518 | img = noise 519 | else: 520 | img = th.randn(*shape, device=device) 521 | indices = list(range(self.num_timesteps))[::-1] 522 | 523 | if progress: 524 | # Lazy import so that we don't depend on tqdm. 525 | from tqdm.auto import tqdm 526 | 527 | indices = tqdm(indices) 528 | 529 | for i in indices: 530 | t = th.tensor([i] * shape[0], device=device) 531 | with th.no_grad(): 532 | out = self.p_sample( 533 | model, 534 | img, 535 | t, 536 | clip_denoised=clip_denoised, 537 | denoised_fn=denoised_fn, 538 | cond_fn=cond_fn, 539 | model_kwargs=model_kwargs, 540 | ) 541 | yield out 542 | img = out["sample"] 543 | 544 | def ddim_sample( 545 | self, 546 | model, 547 | x, 548 | t, 549 | clip_denoised=True, 550 | denoised_fn=None, 551 | cond_fn=None, 552 | model_kwargs=None, 553 | eta=0.0, 554 | ): 555 | """ 556 | Sample x_{t-1} from the model using DDIM. 557 | 558 | Same usage as p_sample(). 559 | """ 560 | out = self.p_mean_variance( 561 | model, 562 | x, 563 | t, 564 | clip_denoised=clip_denoised, 565 | denoised_fn=denoised_fn, 566 | model_kwargs=model_kwargs, 567 | ) 568 | if cond_fn is not None: 569 | out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs) 570 | 571 | # Usually our model outputs epsilon, but we re-derive it 572 | # in case we used x_start or x_prev prediction. 573 | eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"]) 574 | 575 | alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape) 576 | alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape) 577 | sigma = ( 578 | eta 579 | * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) 580 | * th.sqrt(1 - alpha_bar / alpha_bar_prev) 581 | ) 582 | # Equation 12. 583 | noise = th.randn_like(x) 584 | mean_pred = ( 585 | out["pred_xstart"] * th.sqrt(alpha_bar_prev) 586 | + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps 587 | ) 588 | nonzero_mask = ( 589 | (t != 0).float().view(-1, *([1] * (len(x.shape) - 1))) 590 | ) # no noise when t == 0 591 | sample = mean_pred + nonzero_mask * sigma * noise 592 | return {"sample": sample, "pred_xstart": out["pred_xstart"]} 593 | 594 | def ddim_reverse_sample( 595 | self, 596 | model, 597 | x, 598 | t, 599 | clip_denoised=True, 600 | denoised_fn=None, 601 | model_kwargs=None, 602 | eta=0.0, 603 | ): 604 | """ 605 | Sample x_{t+1} from the model using DDIM reverse ODE. 606 | """ 607 | assert eta == 0.0, "Reverse ODE only for deterministic path" 608 | out = self.p_mean_variance( 609 | model, 610 | x, 611 | t, 612 | clip_denoised=clip_denoised, 613 | denoised_fn=denoised_fn, 614 | model_kwargs=model_kwargs, 615 | ) 616 | # Usually our model outputs epsilon, but we re-derive it 617 | # in case we used x_start or x_prev prediction. 618 | eps = ( 619 | _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x 620 | - out["pred_xstart"] 621 | ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape) 622 | alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape) 623 | 624 | # Equation 12. reversed 625 | mean_pred = ( 626 | out["pred_xstart"] * th.sqrt(alpha_bar_next) 627 | + th.sqrt(1 - alpha_bar_next) * eps 628 | ) 629 | 630 | return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]} 631 | 632 | def ddim_sample_loop( 633 | self, 634 | model, 635 | shape, 636 | noise=None, 637 | clip_denoised=True, 638 | denoised_fn=None, 639 | cond_fn=None, 640 | model_kwargs=None, 641 | device=None, 642 | progress=False, 643 | eta=0.0, 644 | ): 645 | """ 646 | Generate samples from the model using DDIM. 647 | 648 | Same usage as p_sample_loop(). 649 | """ 650 | final = None 651 | for sample in self.ddim_sample_loop_progressive( 652 | model, 653 | shape, 654 | noise=noise, 655 | clip_denoised=clip_denoised, 656 | denoised_fn=denoised_fn, 657 | cond_fn=cond_fn, 658 | model_kwargs=model_kwargs, 659 | device=device, 660 | progress=progress, 661 | eta=eta, 662 | ): 663 | final = sample 664 | return final["sample"] 665 | 666 | def ddim_sample_loop_progressive( 667 | self, 668 | model, 669 | shape, 670 | noise=None, 671 | clip_denoised=True, 672 | denoised_fn=None, 673 | cond_fn=None, 674 | model_kwargs=None, 675 | device=None, 676 | progress=False, 677 | eta=0.0, 678 | ): 679 | """ 680 | Use DDIM to sample from the model and yield intermediate samples from 681 | each timestep of DDIM. 682 | 683 | Same usage as p_sample_loop_progressive(). 684 | """ 685 | if device is None: 686 | device = next(model.parameters()).device 687 | assert isinstance(shape, (tuple, list)) 688 | if noise is not None: 689 | img = noise 690 | else: 691 | img = th.randn(*shape, device=device) 692 | indices = list(range(self.num_timesteps))[::-1] 693 | 694 | if progress: 695 | # Lazy import so that we don't depend on tqdm. 696 | from tqdm.auto import tqdm 697 | 698 | indices = tqdm(indices) 699 | 700 | for i in indices: 701 | t = th.tensor([i] * shape[0], device=device) 702 | with th.no_grad(): 703 | out = self.ddim_sample( 704 | model, 705 | img, 706 | t, 707 | clip_denoised=clip_denoised, 708 | denoised_fn=denoised_fn, 709 | cond_fn=cond_fn, 710 | model_kwargs=model_kwargs, 711 | eta=eta, 712 | ) 713 | yield out 714 | img = out["sample"] 715 | 716 | def _vb_terms_bpd( 717 | self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None 718 | ): 719 | """ 720 | Get a term for the variational lower-bound. 721 | 722 | The resulting units are bits (rather than nats, as one might expect). 723 | This allows for comparison to other papers. 724 | 725 | :return: a dict with the following keys: 726 | - 'output': a shape [N] tensor of NLLs or KLs. 727 | - 'pred_xstart': the x_0 predictions. 728 | """ 729 | true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance( 730 | x_start=x_start, x_t=x_t, t=t 731 | ) 732 | out = self.p_mean_variance( 733 | model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs 734 | ) 735 | kl = normal_kl( 736 | true_mean, true_log_variance_clipped, out["mean"], out["log_variance"] 737 | ) 738 | kl = mean_flat(kl) / np.log(2.0) 739 | 740 | decoder_nll = -discretized_gaussian_log_likelihood( 741 | x_start, means=out["mean"], log_scales=0.5 * out["log_variance"] 742 | ) 743 | assert decoder_nll.shape == x_start.shape 744 | decoder_nll = mean_flat(decoder_nll) / np.log(2.0) 745 | 746 | # At the first timestep return the decoder NLL, 747 | # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t)) 748 | output = th.where((t == 0), decoder_nll, kl) 749 | return {"output": output, "pred_xstart": out["pred_xstart"]} 750 | 751 | def training_losses(self, model, x_start, t, model_kwargs=None, noise=None): 752 | """ 753 | Compute training losses for a single timestep. 754 | 755 | :param model: the model to evaluate loss on. 756 | :param x_start: the [N x C x ...] tensor of inputs. 757 | :param t: a batch of timestep indices. 758 | :param model_kwargs: if not None, a dict of extra keyword arguments to 759 | pass to the model. This can be used for conditioning. 760 | :param noise: if specified, the specific Gaussian noise to try to remove. 761 | :return: a dict with the key "loss" containing a tensor of shape [N]. 762 | Some mean or variance settings may also have other keys. 763 | """ 764 | if model_kwargs is None: 765 | model_kwargs = {} 766 | if noise is None: 767 | noise = th.randn_like(x_start) 768 | x_t = self.q_sample(x_start, t, noise=noise) 769 | 770 | terms = {} 771 | 772 | if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL: 773 | terms["loss"] = self._vb_terms_bpd( 774 | model=model, 775 | x_start=x_start, 776 | x_t=x_t, 777 | t=t, 778 | clip_denoised=False, 779 | model_kwargs=model_kwargs, 780 | )["output"] 781 | if self.loss_type == LossType.RESCALED_KL: 782 | terms["loss"] *= self.num_timesteps 783 | elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE: 784 | model_output = model(x_t, self._scale_timesteps(t), **model_kwargs) 785 | 786 | if self.model_var_type in [ 787 | ModelVarType.LEARNED, 788 | ModelVarType.LEARNED_RANGE, 789 | ]: 790 | B, C = x_t.shape[:2] 791 | assert model_output.shape == (B, C * 2, *x_t.shape[2:]) 792 | model_output, model_var_values = th.split(model_output, C, dim=1) 793 | # Learn the variance using the variational bound, but don't let 794 | # it affect our mean prediction. 795 | frozen_out = th.cat([model_output.detach(), model_var_values], dim=1) 796 | terms["vb"] = self._vb_terms_bpd( 797 | model=lambda *args, r=frozen_out: r, 798 | x_start=x_start, 799 | x_t=x_t, 800 | t=t, 801 | clip_denoised=False, 802 | )["output"] 803 | if self.loss_type == LossType.RESCALED_MSE: 804 | # Divide by 1000 for equivalence with initial implementation. 805 | # Without a factor of 1/1000, the VB term hurts the MSE term. 806 | terms["vb"] *= self.num_timesteps / 1000.0 807 | 808 | target = { 809 | ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance( 810 | x_start=x_start, x_t=x_t, t=t 811 | )[0], 812 | ModelMeanType.START_X: x_start, 813 | ModelMeanType.EPSILON: noise, 814 | }[self.model_mean_type] 815 | assert model_output.shape == target.shape == x_start.shape 816 | 817 | # P2 weighting 818 | weight = _extract_into_tensor(1 / (self.p2_k + self.snr)**self.p2_gamma, t, target.shape) 819 | terms["mse"] = mean_flat(weight * (target - model_output) ** 2) 820 | 821 | if "vb" in terms: 822 | terms["loss"] = terms["mse"] + terms["vb"] 823 | else: 824 | terms["loss"] = terms["mse"] 825 | else: 826 | raise NotImplementedError(self.loss_type) 827 | 828 | return terms 829 | 830 | def _prior_bpd(self, x_start): 831 | """ 832 | Get the prior KL term for the variational lower-bound, measured in 833 | bits-per-dim. 834 | 835 | This term can't be optimized, as it only depends on the encoder. 836 | 837 | :param x_start: the [N x C x ...] tensor of inputs. 838 | :return: a batch of [N] KL values (in bits), one per batch element. 839 | """ 840 | batch_size = x_start.shape[0] 841 | t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) 842 | qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) 843 | kl_prior = normal_kl( 844 | mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0 845 | ) 846 | return mean_flat(kl_prior) / np.log(2.0) 847 | 848 | def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None): 849 | """ 850 | Compute the entire variational lower-bound, measured in bits-per-dim, 851 | as well as other related quantities. 852 | 853 | :param model: the model to evaluate loss on. 854 | :param x_start: the [N x C x ...] tensor of inputs. 855 | :param clip_denoised: if True, clip denoised samples. 856 | :param model_kwargs: if not None, a dict of extra keyword arguments to 857 | pass to the model. This can be used for conditioning. 858 | 859 | :return: a dict containing the following keys: 860 | - total_bpd: the total variational lower-bound, per batch element. 861 | - prior_bpd: the prior term in the lower-bound. 862 | - vb: an [N x T] tensor of terms in the lower-bound. 863 | - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep. 864 | - mse: an [N x T] tensor of epsilon MSEs for each timestep. 865 | """ 866 | device = x_start.device 867 | batch_size = x_start.shape[0] 868 | 869 | vb = [] 870 | xstart_mse = [] 871 | mse = [] 872 | for t in list(range(self.num_timesteps))[::-1]: 873 | t_batch = th.tensor([t] * batch_size, device=device) 874 | noise = th.randn_like(x_start) 875 | x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise) 876 | # Calculate VLB term at the current timestep 877 | with th.no_grad(): 878 | out = self._vb_terms_bpd( 879 | model, 880 | x_start=x_start, 881 | x_t=x_t, 882 | t=t_batch, 883 | clip_denoised=clip_denoised, 884 | model_kwargs=model_kwargs, 885 | ) 886 | vb.append(out["output"]) 887 | xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2)) 888 | eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"]) 889 | mse.append(mean_flat((eps - noise) ** 2)) 890 | 891 | vb = th.stack(vb, dim=1) 892 | xstart_mse = th.stack(xstart_mse, dim=1) 893 | mse = th.stack(mse, dim=1) 894 | 895 | prior_bpd = self._prior_bpd(x_start) 896 | total_bpd = vb.sum(dim=1) + prior_bpd 897 | return { 898 | "total_bpd": total_bpd, 899 | "prior_bpd": prior_bpd, 900 | "vb": vb, 901 | "xstart_mse": xstart_mse, 902 | "mse": mse, 903 | } 904 | 905 | 906 | def _extract_into_tensor(arr, timesteps, broadcast_shape): 907 | """ 908 | Extract values from a 1-D numpy array for a batch of indices. 909 | 910 | :param arr: the 1-D numpy array. 911 | :param timesteps: a tensor of indices into the array to extract. 912 | :param broadcast_shape: a larger shape of K dimensions with the batch 913 | dimension equal to the length of timesteps. 914 | :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims. 915 | """ 916 | res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float() 917 | while len(res.shape) < len(broadcast_shape): 918 | res = res[..., None] 919 | return res.expand(broadcast_shape) 920 | --------------------------------------------------------------------------------