├── BCM ├── evaluations │ ├── __init__.py │ ├── requirements.txt │ ├── th_evaluator.py │ └── inception_v3.py ├── cm │ ├── __init__.py │ ├── dist_util.py │ ├── losses.py │ ├── nn.py │ ├── random_util.py │ ├── resample.py │ ├── image_datasets.py │ ├── script_util.py │ ├── fp16_util.py │ └── logger.py ├── docker │ ├── Makefile │ └── Dockerfile ├── scripts │ ├── visualize_image.py │ ├── imagnet64_sample.sh │ ├── bcf_imagenet64_no32_qkv_4096.sh │ ├── cm_train.py │ └── image_sample.py └── setup.py ├── iCT ├── evaluations │ ├── __init__.py │ ├── requirements.txt │ ├── th_evaluator.py │ └── inception_v3.py ├── cm │ ├── __init__.py │ ├── dist_util.py │ ├── losses.py │ ├── nn.py │ ├── image_datasets.py │ ├── random_util.py │ ├── resample.py │ ├── script_util.py │ └── fp16_util.py ├── docker │ ├── Makefile │ └── Dockerfile ├── scripts │ ├── visualize_image.py │ ├── imagenet64_sample.sh │ ├── ict_imagenet64.sh │ ├── ict_imagenet64_no32_qkv_4096.sh │ ├── image_sample.py │ └── cm_train.py └── setup.py ├── LICENSE ├── datasets └── README.md └── README.md /BCM/evaluations/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /iCT/evaluations/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /BCM/evaluations/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow-gpu>=2.0 2 | scipy 3 | requests 4 | tqdm -------------------------------------------------------------------------------- /iCT/evaluations/requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow-gpu>=2.0 2 | scipy 3 | requests 4 | tqdm -------------------------------------------------------------------------------- /BCM/cm/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /iCT/cm/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /BCM/docker/Makefile: -------------------------------------------------------------------------------- 1 | NAME=consistency_models 2 | TAG=0.1 3 | PROJECT_DIRECTORY = $(shell pwd)/.. 4 | 5 | build: 6 | docker build -t ${NAME}:${TAG} -f Dockerfile . 7 | 8 | run: 9 | docker container run --gpus all\ 10 | --restart=always\ 11 | -it -d \ 12 | -v $(PROJECT_DIRECTORY):/home/${NAME}\ 13 | --name ${NAME} ${NAME}:${TAG} /bin/bash 14 | -------------------------------------------------------------------------------- /iCT/docker/Makefile: -------------------------------------------------------------------------------- 1 | NAME=consistency_models 2 | TAG=0.1 3 | PROJECT_DIRECTORY = $(shell pwd)/.. 4 | 5 | build: 6 | docker build -t ${NAME}:${TAG} -f Dockerfile . 7 | 8 | run: 9 | docker container run --gpus all\ 10 | --restart=always\ 11 | -it -d \ 12 | -v $(PROJECT_DIRECTORY):/home/${NAME}\ 13 | --name ${NAME} ${NAME}:${TAG} /bin/bash 14 | -------------------------------------------------------------------------------- /BCM/scripts/visualize_image.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | 4 | 5 | vis_num = 64 6 | npz_path = 'PATH_TO_THE_SAMPLE_NPZ_FILE' 7 | save_path = f'{npz_path[:-4]}.png' 8 | img = np.load(npz_path)['arr_0'][:vis_num] 9 | 10 | samples = img.reshape((8, 8, 64, 64, 3)) 11 | samples = samples.transpose((0, 2, 1, 3, 4)) 12 | samples = samples.reshape( 13 | (samples.shape[0] * samples.shape[1], samples.shape[2] * samples.shape[3], samples.shape[4])) 14 | plt.figure(figsize=(10, 10)) 15 | plt.imshow(samples) 16 | plt.axis('off') 17 | plt.savefig(save_path) 18 | plt.close() 19 | 20 | -------------------------------------------------------------------------------- /iCT/scripts/visualize_image.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | 4 | 5 | vis_num = 64 6 | npz_path = 'PATH_TO_THE_SAMPLE_NPZ_FILE' 7 | save_path = f'{npz_path[:-4]}.png' 8 | img = np.load(npz_path)['arr_0'][:vis_num] 9 | 10 | samples = img.reshape((8, 8, 64, 64, 3)) 11 | samples = samples.transpose((0, 2, 1, 3, 4)) 12 | samples = samples.reshape( 13 | (samples.shape[0] * samples.shape[1], samples.shape[2] * samples.shape[3], samples.shape[4])) 14 | plt.figure(figsize=(10, 10)) 15 | plt.imshow(samples) 16 | plt.axis('off') 17 | plt.savefig(save_path) 18 | plt.close() 19 | 20 | -------------------------------------------------------------------------------- /iCT/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="consistency-models", 5 | py_modules=["cm", "evaluations"], 6 | install_requires=[ 7 | "blobfile>=1.0.5", 8 | "torch", 9 | "tqdm", 10 | "numpy", 11 | "scipy", 12 | "pandas", 13 | "Cython", 14 | "piq==0.7.0", 15 | "joblib==0.14.0", 16 | "albumentations==0.4.3", 17 | "lmdb", 18 | "clip @ git+https://github.com/openai/CLIP.git", 19 | "mpi4py", 20 | "flash-attn==0.2.8", # optional 21 | "pillow", 22 | ], 23 | ) 24 | -------------------------------------------------------------------------------- /BCM/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="consistency-models", 5 | py_modules=["cm", "evaluations"], 6 | install_requires=[ 7 | "blobfile>=1.0.5", 8 | "torch", 9 | "tqdm", 10 | "numpy", 11 | "scipy", 12 | "pandas", 13 | "Cython", 14 | "piq==0.7.0", 15 | "joblib==0.14.0", 16 | "albumentations==0.4.3", 17 | "lmdb", 18 | "clip @ git+https://github.com/openai/CLIP.git", 19 | "mpi4py", 20 | # "flash-attn==0.2.8", # optional 21 | "pillow", 22 | ], 23 | ) 24 | -------------------------------------------------------------------------------- /iCT/scripts/imagenet64_sample.sh: -------------------------------------------------------------------------------- 1 | python -u scripts/image_sample.py \ 2 | --batch_size 32 \ 3 | --training_mode consistency_training \ 4 | --sampler onestep \ 5 | --model_path CKPT_DIR/ict_imagenet64_no32_qkv_4096/ema_0.99997_680000.pt \ 6 | --save_dir samples \ 7 | --exp_name ict_imagenet64_no32_qkv_4096 \ 8 | --attention_resolutions 16,8 \ 9 | --class_cond True \ 10 | --use_scale_shift_norm False \ 11 | --dropout 0.0 \ 12 | --image_size 64 \ 13 | --num_channels 192 \ 14 | --num_head_channels 64 \ 15 | --num_res_blocks 3 \ 16 | --num_samples 50000 \ 17 | --resblock_updown True \ 18 | --use_fp16 True \ 19 | --weight_schedule uniform -------------------------------------------------------------------------------- /BCM/scripts/imagnet64_sample.sh: -------------------------------------------------------------------------------- 1 | python -u scripts/image_sample.py \ 2 | --batch_size 32 \ 3 | --training_mode consistency_training \ 4 | --sampler onestep \ 5 | --model_path CKPT_DIR/bcf_imagenet64_no32_qkv_4096/ema_0.99997_234000.pt \ 6 | --save_dir samples \ 7 | --exp_name bcf_imagenet64_no32_qkv_4096 \ 8 | --attention_resolutions 16,8 \ 9 | --class_cond True \ 10 | --use_scale_shift_norm False \ 11 | --dropout 0.0 \ 12 | --image_size 64 \ 13 | --num_channels 192 \ 14 | --num_head_channels 64 \ 15 | --num_res_blocks 3 \ 16 | --num_samples 50000 \ 17 | --resblock_updown True \ 18 | --use_fp16 True \ 19 | --weight_schedule uniform \ 20 | --eval_mse False \ 21 | --test_data_dir IMAGENET_DIR/val -------------------------------------------------------------------------------- /BCM/docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 2 | 3 | ENV DEBIAN_FRONTEND=noninteractive PIP_PREFER_BINARY=1 4 | 5 | RUN apt-get update && apt-get install -y --no-install-recommends \ 6 | libgl1-mesa-dev libopenmpi-dev git wget \ 7 | python3 python3-dev python3-pip python3-setuptools python3-wheel \ 8 | && apt-get clean && rm -rf /var/lib/apt/lists/* 9 | 10 | RUN echo "export PATH=/usr/local/cuda/bin:$PATH" >> /etc/bash.bashrc \ 11 | && echo "export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH" >> /etc/bash.bashrc 12 | 13 | RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel packaging mpi4py \ 14 | && pip3 install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu118 \ 15 | && pip3 install flash-attn==0.2.8 16 | 17 | WORKDIR /home/ 18 | RUN pip3 install -e git+https://github.com/openai/consistency_models.git@main#egg=consistency_models \ 19 | && ln -s /usr/bin/python3 /usr/bin/python 20 | -------------------------------------------------------------------------------- /iCT/docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 2 | 3 | ENV DEBIAN_FRONTEND=noninteractive PIP_PREFER_BINARY=1 4 | 5 | RUN apt-get update && apt-get install -y --no-install-recommends \ 6 | libgl1-mesa-dev libopenmpi-dev git wget \ 7 | python3 python3-dev python3-pip python3-setuptools python3-wheel \ 8 | && apt-get clean && rm -rf /var/lib/apt/lists/* 9 | 10 | RUN echo "export PATH=/usr/local/cuda/bin:$PATH" >> /etc/bash.bashrc \ 11 | && echo "export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH" >> /etc/bash.bashrc 12 | 13 | RUN pip3 install --no-cache-dir --upgrade pip setuptools wheel packaging mpi4py \ 14 | && pip3 install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu118 \ 15 | && pip3 install flash-attn==0.2.8 16 | 17 | WORKDIR /home/ 18 | RUN pip3 install -e git+https://github.com/openai/consistency_models.git@main#egg=consistency_models \ 19 | && ln -s /usr/bin/python3 /usr/bin/python 20 | -------------------------------------------------------------------------------- /iCT/scripts/ict_imagenet64.sh: -------------------------------------------------------------------------------- 1 | python -u scripts/cm_train.py \ 2 | --training_mode consistency_training \ 3 | --exp_name ict_imagenet64 \ 4 | --target_ema_mode adaptive \ 5 | --start_ema 0.95 \ 6 | --log_interval 100 \ 7 | --save_interval 20000 \ 8 | --scale_mode progressive \ 9 | --start_scales 10 \ 10 | --end_scales 1280 \ 11 | --total_training_steps 800000 \ 12 | --loss_norm ict \ 13 | --lr_anneal_steps 0 \ 14 | --attention_resolutions 32,16,8 \ 15 | --class_cond True \ 16 | --use_scale_shift_norm False \ 17 | --dropout 0.0 \ 18 | --teacher_dropout 0.1 \ 19 | --ema_rate 0.99997 \ 20 | --global_batch_size 4096 \ 21 | --image_size 64 \ 22 | --lr 0.0001 \ 23 | --num_channels 192 \ 24 | --num_head_channels 64 \ 25 | --num_res_blocks 3 \ 26 | --resblock_updown True \ 27 | --schedule_sampler uniform \ 28 | --use_fp16 True \ 29 | --weight_decay 0.0 \ 30 | --weight_schedule ict \ 31 | --data_dir IMAGENET_PATH/train \ 32 | # --resume_checkpoint CKPT_DIR/model500000.pt \ -------------------------------------------------------------------------------- /iCT/scripts/ict_imagenet64_no32_qkv_4096.sh: -------------------------------------------------------------------------------- 1 | python -u scripts/cm_train.py \ 2 | --training_mode consistency_training \ 3 | --exp_name ict_imagenet64_no32_qkv_4096 \ 4 | --target_ema_mode adaptive \ 5 | --start_ema 0.95 \ 6 | --log_interval 100 \ 7 | --save_interval 20000 \ 8 | --scale_mode progressive \ 9 | --start_scales 10 \ 10 | --end_scales 1280 \ 11 | --total_training_steps 800000 \ 12 | --loss_norm ict \ 13 | --lr_anneal_steps 0 \ 14 | --attention_resolutions 16,8 \ 15 | --class_cond True \ 16 | --use_scale_shift_norm False \ 17 | --dropout 0.0 \ 18 | --teacher_dropout 0.1 \ 19 | --ema_rate 0.99997 \ 20 | --global_batch_size 4096 \ 21 | --image_size 64 \ 22 | --lr 0.0001 \ 23 | --num_channels 192 \ 24 | --num_head_channels 64 \ 25 | --num_res_blocks 3 \ 26 | --resblock_updown True \ 27 | --schedule_sampler uniform \ 28 | --use_fp16 True \ 29 | --weight_decay 0.0 \ 30 | --weight_schedule ict \ 31 | --data_dir IMAGENET_PATH/train \ 32 | # --resume_checkpoint CKPT_DIR/model500000.pt \ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Liangchen Li and Jiajun He 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. 22 | -------------------------------------------------------------------------------- /BCM/scripts/bcf_imagenet64_no32_qkv_4096.sh: -------------------------------------------------------------------------------- 1 | python -u scripts/cm_train.py \ 2 | --training_mode consistency_training \ 3 | --exp_name bcf_imagenet64_no32_qkv_4096 \ 4 | --target_ema_mode adaptive \ 5 | --start_ema 0.95 \ 6 | --log_interval 100 \ 7 | --save_interval 10000 \ 8 | --scale_mode progressive \ 9 | --start_scales 320 \ 10 | --end_scales 320 \ 11 | --total_training_steps 800000 \ 12 | --loss_norm ict \ 13 | --lr_anneal_steps 0 \ 14 | --attention_resolutions 16,8 \ 15 | --class_cond True \ 16 | --use_scale_shift_norm False \ 17 | --dropout 0.0 \ 18 | --teacher_dropout 0.1 \ 19 | --ema_rate 0.99997 \ 20 | --global_batch_size 4096 \ 21 | --image_size 64 \ 22 | --lr 0.0001 \ 23 | --num_channels 192 \ 24 | --num_head_channels 64 \ 25 | --num_res_blocks 3 \ 26 | --resblock_updown True \ 27 | --schedule_sampler uniform \ 28 | --use_fp16 True \ 29 | --weight_decay 0.0 \ 30 | --weight_schedule ict \ 31 | --data_dir IMAGENET_PATH/train \ 32 | --bcf True \ 33 | --pretrained_model_path CKPT_DIR/ict_imagenet64_no32_qkv_4096/ema_0.99997_600000.pt \ 34 | # --resume_checkpoint CKPT_DIR/bcf_imagenet64_no32_qkv_4096/model234000.pt -------------------------------------------------------------------------------- /datasets/README.md: -------------------------------------------------------------------------------- 1 | # Downloading datasets 2 | 3 | This directory includes instructions and scripts for downloading ImageNet 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](https://image-net.org/challenges/LSVRC/2012/2012-downloads.php) 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 | -------------------------------------------------------------------------------- /BCM/cm/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("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 | -------------------------------------------------------------------------------- /iCT/cm/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("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 | -------------------------------------------------------------------------------- /BCM/cm/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 | -------------------------------------------------------------------------------- /iCT/cm/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 | -------------------------------------------------------------------------------- /iCT/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 | import sys 6 | sys.path.append('/mnt/petrelfs/liliangchen/try/consistency_model_torch/cm_v2') 7 | 8 | import argparse 9 | import os 10 | 11 | import numpy as np 12 | import torch as th 13 | import torch.distributed as dist 14 | 15 | from cm import dist_util, logger 16 | from cm.script_util import ( 17 | NUM_CLASSES, 18 | model_and_diffusion_defaults, 19 | create_model_and_diffusion, 20 | add_dict_to_argparser, 21 | args_to_dict, 22 | ) 23 | from cm.random_util import get_generator 24 | from cm.karras_diffusion import karras_sample 25 | 26 | 27 | def main(): 28 | args = create_argparser().parse_args() 29 | 30 | dist_util.setup_dist() 31 | logger.configure(dir=os.path.join(args.save_dir, args.exp_name)) 32 | 33 | if "consistency" in args.training_mode: 34 | distillation = True 35 | else: 36 | distillation = False 37 | 38 | logger.log("creating model and diffusion...") 39 | model, diffusion = create_model_and_diffusion( 40 | **args_to_dict(args, model_and_diffusion_defaults().keys()), 41 | distillation=distillation, 42 | ) 43 | model.load_state_dict( 44 | dist_util.load_state_dict(args.model_path, map_location="cpu") 45 | ) 46 | model.to(dist_util.dev()) 47 | if args.use_fp16: 48 | model.convert_to_fp16() 49 | model.eval() 50 | model_path = args.model_path 51 | model_path = model_path.split('/')[-1][:-3] 52 | 53 | logger.log("sampling...") 54 | 55 | all_images = [] 56 | all_labels = [] 57 | 58 | while len(all_images) * args.batch_size < args.num_samples: 59 | model_kwargs = {} 60 | if args.class_cond: 61 | classes = th.randint( 62 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 63 | ) 64 | model_kwargs["y"] = classes 65 | 66 | sample = th.randn((args.batch_size, 3, args.image_size, args.image_size), device=dist_util.dev()) * 80.0 67 | multiplier = th.ones(sample.shape[0], dtype=sample.dtype, device=sample.device) 68 | 69 | # one-step sampling 70 | _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 80.0, y=classes) 71 | 72 | # uncomment the lines below for two-step sampling 73 | # sample += th.randn_like(sample, device=sample.device) * 0.7 74 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 0.7, y=classes) 75 | 76 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 77 | sample = sample.permute(0, 2, 3, 1) 78 | sample = sample.contiguous() 79 | 80 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 81 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 82 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 83 | if args.class_cond: 84 | gathered_labels = [ 85 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 86 | ] 87 | dist.all_gather(gathered_labels, classes) 88 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 89 | logger.log(f"created {len(all_images) * args.batch_size} samples") 90 | 91 | arr = np.concatenate(all_images, axis=0) 92 | arr = arr[: args.num_samples] 93 | if args.class_cond: 94 | label_arr = np.concatenate(all_labels, axis=0) 95 | label_arr = label_arr[: args.num_samples] 96 | if dist.get_rank() == 0: 97 | shape_str = "x".join([str(x) for x in arr.shape]) 98 | out_path = os.path.join(logger.get_dir(), f"{model_path}_samples_{shape_str}.npz") 99 | logger.log(f"saving to {out_path}") 100 | if args.class_cond: 101 | np.savez(out_path, arr, label_arr) 102 | else: 103 | np.savez(out_path, arr) 104 | 105 | dist.barrier() 106 | logger.log("sampling complete") 107 | 108 | 109 | def create_argparser(): 110 | defaults = dict( 111 | training_mode="edm", 112 | generator="determ", 113 | clip_denoised=True, 114 | num_samples=50000, 115 | batch_size=16, 116 | sampler="heun", 117 | s_churn=0.0, 118 | s_tmin=0.0, 119 | s_tmax=float("inf"), 120 | s_noise=1.0, 121 | save_dir='./checkpoints', 122 | exp_name='ict', 123 | steps=40, 124 | model_path="", 125 | seed=42, 126 | ts="", 127 | ) 128 | defaults.update(model_and_diffusion_defaults()) 129 | parser = argparse.ArgumentParser() 130 | add_dict_to_argparser(parser, defaults) 131 | return parser 132 | 133 | 134 | if __name__ == "__main__": 135 | main() 136 | -------------------------------------------------------------------------------- /iCT/scripts/cm_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | import sys 5 | sys.path.append('../iCT/') 6 | 7 | import argparse 8 | import os 9 | 10 | from cm import dist_util, logger 11 | from cm.image_datasets import load_data 12 | from cm.resample import create_named_schedule_sampler 13 | from cm.script_util import ( 14 | model_and_diffusion_defaults, 15 | create_model_and_diffusion, 16 | cm_train_defaults, 17 | args_to_dict, 18 | add_dict_to_argparser, 19 | create_ema_and_scales_fn, 20 | ) 21 | from cm.train_util import CMTrainLoop 22 | import torch.distributed as dist 23 | import copy 24 | 25 | 26 | def main(): 27 | args = create_argparser().parse_args() 28 | 29 | dist_util.setup_dist() 30 | logger.configure(dir=os.path.join(args.save_dir, args.exp_name)) 31 | 32 | logger.log("creating model and diffusion...") 33 | ema_scale_fn = create_ema_and_scales_fn( 34 | target_ema_mode=args.target_ema_mode, 35 | start_ema=args.start_ema, 36 | scale_mode=args.scale_mode, 37 | start_scales=args.start_scales, 38 | end_scales=args.end_scales, 39 | total_steps=args.total_training_steps, 40 | distill_steps_per_iter=args.distill_steps_per_iter, 41 | ) 42 | if args.training_mode == "progdist": 43 | distillation = False 44 | elif "consistency" in args.training_mode: 45 | distillation = True 46 | else: 47 | raise ValueError(f"unknown training mode {args.training_mode}") 48 | 49 | model_and_diffusion_kwargs = args_to_dict( 50 | args, model_and_diffusion_defaults().keys() 51 | ) 52 | model_and_diffusion_kwargs["distillation"] = distillation 53 | model, diffusion = create_model_and_diffusion(**model_and_diffusion_kwargs) 54 | model.to(dist_util.dev()) 55 | model.train() 56 | if args.use_fp16: 57 | model.convert_to_fp16() 58 | 59 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 60 | 61 | logger.log("creating data loader...") 62 | if args.batch_size == -1: 63 | batch_size = args.global_batch_size // dist.get_world_size() 64 | if args.global_batch_size % dist.get_world_size() != 0: 65 | logger.log( 66 | f"warning, using smaller global_batch_size of {dist.get_world_size()*batch_size} instead of {args.global_batch_size}" 67 | ) 68 | else: 69 | batch_size = args.batch_size 70 | 71 | data = load_data( 72 | data_dir=args.data_dir, 73 | batch_size=batch_size, 74 | image_size=args.image_size, 75 | class_cond=args.class_cond, 76 | ) 77 | 78 | if len(args.teacher_model_path) > 0: # path to the teacher score model. 79 | logger.log(f"loading the teacher model from {args.teacher_model_path}") 80 | teacher_model_and_diffusion_kwargs = copy.deepcopy(model_and_diffusion_kwargs) 81 | teacher_model_and_diffusion_kwargs["dropout"] = args.teacher_dropout 82 | teacher_model_and_diffusion_kwargs["distillation"] = False 83 | teacher_model, teacher_diffusion = create_model_and_diffusion( 84 | **teacher_model_and_diffusion_kwargs, 85 | ) 86 | 87 | teacher_model.load_state_dict( 88 | dist_util.load_state_dict(args.teacher_model_path, map_location="cpu"), 89 | ) 90 | 91 | teacher_model.to(dist_util.dev()) 92 | teacher_model.eval() 93 | 94 | for dst, src in zip(model.parameters(), teacher_model.parameters()): 95 | dst.data.copy_(src.data) 96 | 97 | if args.use_fp16: 98 | teacher_model.convert_to_fp16() 99 | 100 | else: 101 | teacher_model = None 102 | teacher_diffusion = None 103 | 104 | # load the target model for distillation, if path specified. 105 | 106 | logger.log("creating the target model") 107 | target_model, _ = create_model_and_diffusion( 108 | **model_and_diffusion_kwargs, 109 | ) 110 | 111 | target_model.to(dist_util.dev()) 112 | target_model.train() 113 | 114 | dist_util.sync_params(target_model.parameters()) 115 | dist_util.sync_params(target_model.buffers()) 116 | 117 | for dst, src in zip(target_model.parameters(), model.parameters()): 118 | dst.data.copy_(src.data) 119 | 120 | if args.use_fp16: 121 | target_model.convert_to_fp16() 122 | 123 | logger.log("training...") 124 | CMTrainLoop( 125 | model=model, 126 | target_model=target_model, 127 | teacher_model=teacher_model, 128 | teacher_diffusion=teacher_diffusion, 129 | training_mode=args.training_mode, 130 | ema_scale_fn=ema_scale_fn, 131 | total_training_steps=args.total_training_steps, 132 | diffusion=diffusion, 133 | data=data, 134 | batch_size=batch_size, 135 | microbatch=args.microbatch, 136 | lr=args.lr, 137 | ema_rate=args.ema_rate, 138 | log_interval=args.log_interval, 139 | save_interval=args.save_interval, 140 | resume_checkpoint=args.resume_checkpoint, 141 | use_fp16=args.use_fp16, 142 | fp16_scale_growth=args.fp16_scale_growth, 143 | schedule_sampler=schedule_sampler, 144 | weight_decay=args.weight_decay, 145 | lr_anneal_steps=args.lr_anneal_steps, 146 | ).run_loop() 147 | 148 | 149 | def create_argparser(): 150 | defaults = dict( 151 | data_dir="", 152 | schedule_sampler="uniform", 153 | lr=1e-4, 154 | weight_decay=0.0, 155 | lr_anneal_steps=0, 156 | global_batch_size=2048, 157 | batch_size=-1, 158 | microbatch=-1, # -1 disables microbatches 159 | ema_rate="0.9999", # comma-separated list of EMA values 160 | log_interval=10, 161 | save_interval=10000, 162 | save_dir='./checkpoints', 163 | exp_name='ict', 164 | resume_checkpoint="", 165 | use_fp16=False, 166 | fp16_scale_growth=1e-3, 167 | ) 168 | defaults.update(model_and_diffusion_defaults()) 169 | defaults.update(cm_train_defaults()) 170 | parser = argparse.ArgumentParser() 171 | add_dict_to_argparser(parser, defaults) 172 | return parser 173 | 174 | 175 | if __name__ == "__main__": 176 | main() 177 | -------------------------------------------------------------------------------- /BCM/cm/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 | import numpy as np 10 | import torch.nn.functional as F 11 | 12 | 13 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 14 | class SiLU(nn.Module): 15 | def forward(self, x): 16 | return x * th.sigmoid(x) 17 | 18 | 19 | class GroupNorm32(nn.GroupNorm): 20 | def forward(self, x): 21 | return super().forward(x.float()).type(x.dtype) 22 | 23 | 24 | def conv_nd(dims, *args, **kwargs): 25 | """ 26 | Create a 1D, 2D, or 3D convolution module. 27 | """ 28 | if dims == 1: 29 | return nn.Conv1d(*args, **kwargs) 30 | elif dims == 2: 31 | return nn.Conv2d(*args, **kwargs) 32 | elif dims == 3: 33 | return nn.Conv3d(*args, **kwargs) 34 | raise ValueError(f"unsupported dimensions: {dims}") 35 | 36 | 37 | def linear(*args, **kwargs): 38 | """ 39 | Create a linear module. 40 | """ 41 | return nn.Linear(*args, **kwargs) 42 | 43 | 44 | def avg_pool_nd(dims, *args, **kwargs): 45 | """ 46 | Create a 1D, 2D, or 3D average pooling module. 47 | """ 48 | if dims == 1: 49 | return nn.AvgPool1d(*args, **kwargs) 50 | elif dims == 2: 51 | return nn.AvgPool2d(*args, **kwargs) 52 | elif dims == 3: 53 | return nn.AvgPool3d(*args, **kwargs) 54 | raise ValueError(f"unsupported dimensions: {dims}") 55 | 56 | 57 | def update_ema(target_params, source_params, rate=0.99): 58 | """ 59 | Update target parameters to be closer to those of source parameters using 60 | an exponential moving average. 61 | 62 | :param target_params: the target parameter sequence. 63 | :param source_params: the source parameter sequence. 64 | :param rate: the EMA rate (closer to 1 means slower). 65 | """ 66 | for targ, src in zip(target_params, source_params): 67 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 68 | 69 | 70 | def zero_module(module): 71 | """ 72 | Zero out the parameters of a module and return it. 73 | """ 74 | for p in module.parameters(): 75 | p.detach().zero_() 76 | return module 77 | 78 | 79 | def scale_module(module, scale): 80 | """ 81 | Scale the parameters of a module and return it. 82 | """ 83 | for p in module.parameters(): 84 | p.detach().mul_(scale) 85 | return module 86 | 87 | 88 | def mean_flat(tensor): 89 | """ 90 | Take the mean over all non-batch dimensions. 91 | """ 92 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 93 | 94 | 95 | def append_dims(x, target_dims): 96 | """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" 97 | dims_to_append = target_dims - x.ndim 98 | if dims_to_append < 0: 99 | raise ValueError( 100 | f"input has {x.ndim} dims but target_dims is {target_dims}, which is less" 101 | ) 102 | return x[(...,) + (None,) * dims_to_append] 103 | 104 | 105 | def append_zero(x): 106 | return th.cat([x, x.new_zeros([1])]) 107 | 108 | 109 | def normalization(channels): 110 | """ 111 | Make a standard normalization layer. 112 | 113 | :param channels: number of input channels. 114 | :return: an nn.Module for normalization. 115 | """ 116 | return GroupNorm32(32, channels) 117 | 118 | 119 | def timestep_embedding(timesteps, dim, max_period=10000): 120 | """ 121 | Create sinusoidal timestep embeddings. 122 | 123 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 124 | These may be fractional. 125 | :param dim: the dimension of the output. 126 | :param max_period: controls the minimum frequency of the embeddings. 127 | :return: an [N x dim] Tensor of positional embeddings. 128 | """ 129 | half = dim // 2 130 | freqs = th.exp( 131 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 132 | ).to(device=timesteps.device) 133 | args = timesteps[:, None].float() * freqs[None] 134 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 135 | if dim % 2: 136 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 137 | return embedding 138 | 139 | 140 | def checkpoint(func, inputs, params, flag): 141 | """ 142 | Evaluate a function without caching intermediate activations, allowing for 143 | reduced memory at the expense of extra compute in the backward pass. 144 | 145 | :param func: the function to evaluate. 146 | :param inputs: the argument sequence to pass to `func`. 147 | :param params: a sequence of parameters `func` depends on but does not 148 | explicitly take as arguments. 149 | :param flag: if False, disable gradient checkpointing. 150 | """ 151 | if flag: 152 | args = tuple(inputs) + tuple(params) 153 | return CheckpointFunction.apply(func, len(inputs), *args) 154 | else: 155 | return func(*inputs) 156 | 157 | 158 | class CheckpointFunction(th.autograd.Function): 159 | @staticmethod 160 | def forward(ctx, run_function, length, *args): 161 | ctx.run_function = run_function 162 | ctx.input_tensors = list(args[:length]) 163 | ctx.input_params = list(args[length:]) 164 | with th.no_grad(): 165 | output_tensors = ctx.run_function(*ctx.input_tensors) 166 | return output_tensors 167 | 168 | @staticmethod 169 | def backward(ctx, *output_grads): 170 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 171 | with th.enable_grad(): 172 | # Fixes a bug where the first op in run_function modifies the 173 | # Tensor storage in place, which is not allowed for detach()'d 174 | # Tensors. 175 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 176 | output_tensors = ctx.run_function(*shallow_copies) 177 | input_grads = th.autograd.grad( 178 | output_tensors, 179 | ctx.input_tensors + ctx.input_params, 180 | output_grads, 181 | allow_unused=True, 182 | ) 183 | del ctx.input_tensors 184 | del ctx.input_params 185 | del output_tensors 186 | return (None, None) + input_grads 187 | -------------------------------------------------------------------------------- /iCT/cm/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 | import numpy as np 10 | import torch.nn.functional as F 11 | 12 | 13 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 14 | class SiLU(nn.Module): 15 | def forward(self, x): 16 | return x * th.sigmoid(x) 17 | 18 | 19 | class GroupNorm32(nn.GroupNorm): 20 | def forward(self, x): 21 | return super().forward(x.float()).type(x.dtype) 22 | 23 | 24 | def conv_nd(dims, *args, **kwargs): 25 | """ 26 | Create a 1D, 2D, or 3D convolution module. 27 | """ 28 | if dims == 1: 29 | return nn.Conv1d(*args, **kwargs) 30 | elif dims == 2: 31 | return nn.Conv2d(*args, **kwargs) 32 | elif dims == 3: 33 | return nn.Conv3d(*args, **kwargs) 34 | raise ValueError(f"unsupported dimensions: {dims}") 35 | 36 | 37 | def linear(*args, **kwargs): 38 | """ 39 | Create a linear module. 40 | """ 41 | return nn.Linear(*args, **kwargs) 42 | 43 | 44 | def avg_pool_nd(dims, *args, **kwargs): 45 | """ 46 | Create a 1D, 2D, or 3D average pooling module. 47 | """ 48 | if dims == 1: 49 | return nn.AvgPool1d(*args, **kwargs) 50 | elif dims == 2: 51 | return nn.AvgPool2d(*args, **kwargs) 52 | elif dims == 3: 53 | return nn.AvgPool3d(*args, **kwargs) 54 | raise ValueError(f"unsupported dimensions: {dims}") 55 | 56 | 57 | def update_ema(target_params, source_params, rate=0.99): 58 | """ 59 | Update target parameters to be closer to those of source parameters using 60 | an exponential moving average. 61 | 62 | :param target_params: the target parameter sequence. 63 | :param source_params: the source parameter sequence. 64 | :param rate: the EMA rate (closer to 1 means slower). 65 | """ 66 | for targ, src in zip(target_params, source_params): 67 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 68 | 69 | 70 | def zero_module(module): 71 | """ 72 | Zero out the parameters of a module and return it. 73 | """ 74 | for p in module.parameters(): 75 | p.detach().zero_() 76 | return module 77 | 78 | 79 | def scale_module(module, scale): 80 | """ 81 | Scale the parameters of a module and return it. 82 | """ 83 | for p in module.parameters(): 84 | p.detach().mul_(scale) 85 | return module 86 | 87 | 88 | def mean_flat(tensor): 89 | """ 90 | Take the mean over all non-batch dimensions. 91 | """ 92 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 93 | 94 | 95 | def append_dims(x, target_dims): 96 | """Appends dimensions to the end of a tensor until it has target_dims dimensions.""" 97 | dims_to_append = target_dims - x.ndim 98 | if dims_to_append < 0: 99 | raise ValueError( 100 | f"input has {x.ndim} dims but target_dims is {target_dims}, which is less" 101 | ) 102 | return x[(...,) + (None,) * dims_to_append] 103 | 104 | 105 | def append_zero(x): 106 | return th.cat([x, x.new_zeros([1])]) 107 | 108 | 109 | def normalization(channels): 110 | """ 111 | Make a standard normalization layer. 112 | 113 | :param channels: number of input channels. 114 | :return: an nn.Module for normalization. 115 | """ 116 | return GroupNorm32(32, channels) 117 | 118 | 119 | def timestep_embedding(timesteps, dim, max_period=10000): 120 | """ 121 | Create sinusoidal timestep embeddings. 122 | 123 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 124 | These may be fractional. 125 | :param dim: the dimension of the output. 126 | :param max_period: controls the minimum frequency of the embeddings. 127 | :return: an [N x dim] Tensor of positional embeddings. 128 | """ 129 | half = dim // 2 130 | freqs = th.exp( 131 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 132 | ).to(device=timesteps.device) 133 | args = timesteps[:, None].float() * freqs[None] 134 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 135 | if dim % 2: 136 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 137 | return embedding 138 | 139 | 140 | def checkpoint(func, inputs, params, flag): 141 | """ 142 | Evaluate a function without caching intermediate activations, allowing for 143 | reduced memory at the expense of extra compute in the backward pass. 144 | 145 | :param func: the function to evaluate. 146 | :param inputs: the argument sequence to pass to `func`. 147 | :param params: a sequence of parameters `func` depends on but does not 148 | explicitly take as arguments. 149 | :param flag: if False, disable gradient checkpointing. 150 | """ 151 | if flag: 152 | args = tuple(inputs) + tuple(params) 153 | return CheckpointFunction.apply(func, len(inputs), *args) 154 | else: 155 | return func(*inputs) 156 | 157 | 158 | class CheckpointFunction(th.autograd.Function): 159 | @staticmethod 160 | def forward(ctx, run_function, length, *args): 161 | ctx.run_function = run_function 162 | ctx.input_tensors = list(args[:length]) 163 | ctx.input_params = list(args[length:]) 164 | with th.no_grad(): 165 | output_tensors = ctx.run_function(*ctx.input_tensors) 166 | return output_tensors 167 | 168 | @staticmethod 169 | def backward(ctx, *output_grads): 170 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 171 | with th.enable_grad(): 172 | # Fixes a bug where the first op in run_function modifies the 173 | # Tensor storage in place, which is not allowed for detach()'d 174 | # Tensors. 175 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 176 | output_tensors = ctx.run_function(*shallow_copies) 177 | input_grads = th.autograd.grad( 178 | output_tensors, 179 | ctx.input_tensors + ctx.input_params, 180 | output_grads, 181 | allow_unused=True, 182 | ) 183 | del ctx.input_tensors 184 | del ctx.input_params 185 | del output_tensors 186 | return (None, None) + input_grads 187 | -------------------------------------------------------------------------------- /iCT/cm/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 | -------------------------------------------------------------------------------- /BCM/cm/random_util.py: -------------------------------------------------------------------------------- 1 | import torch as th 2 | import torch.distributed as dist 3 | from . import dist_util 4 | 5 | 6 | def get_generator(generator, num_samples=0, seed=0): 7 | if generator == "dummy": 8 | return DummyGenerator() 9 | elif generator == "determ": 10 | return DeterministicGenerator(num_samples, seed) 11 | elif generator == "determ-indiv": 12 | return DeterministicIndividualGenerator(num_samples, seed) 13 | else: 14 | raise NotImplementedError 15 | 16 | 17 | class DummyGenerator: 18 | def randn(self, *args, **kwargs): 19 | return th.randn(*args, **kwargs) 20 | 21 | def randint(self, *args, **kwargs): 22 | return th.randint(*args, **kwargs) 23 | 24 | def randn_like(self, *args, **kwargs): 25 | return th.randn_like(*args, **kwargs) 26 | 27 | 28 | class DeterministicGenerator: 29 | """ 30 | RNG to deterministically sample num_samples samples that does not depend on batch_size or mpi_machines 31 | Uses a single rng and samples num_samples sized randomness and subsamples the current indices 32 | """ 33 | 34 | def __init__(self, num_samples, seed=0): 35 | if dist.is_initialized(): 36 | self.rank = dist.get_rank() 37 | self.world_size = dist.get_world_size() 38 | else: 39 | print("Warning: Distributed not initialised, using single rank") 40 | self.rank = 0 41 | self.world_size = 1 42 | self.num_samples = num_samples 43 | self.done_samples = 0 44 | self.seed = seed 45 | self.rng_cpu = th.Generator() 46 | if th.cuda.is_available(): 47 | self.rng_cuda = th.Generator(dist_util.dev()) 48 | self.set_seed(seed) 49 | 50 | def get_global_size_and_indices(self, size): 51 | global_size = (self.num_samples, *size[1:]) 52 | indices = th.arange( 53 | self.done_samples + self.rank, 54 | self.done_samples + self.world_size * int(size[0]), 55 | self.world_size, 56 | ) 57 | indices = th.clamp(indices, 0, self.num_samples - 1) 58 | assert ( 59 | len(indices) == size[0] 60 | ), f"rank={self.rank}, ws={self.world_size}, l={len(indices)}, bs={size[0]}" 61 | return global_size, indices 62 | 63 | def get_generator(self, device): 64 | return self.rng_cpu if th.device(device).type == "cpu" else self.rng_cuda 65 | 66 | def randn(self, *size, dtype=th.float, device="cpu"): 67 | global_size, indices = self.get_global_size_and_indices(size) 68 | generator = self.get_generator(device) 69 | return th.randn(*global_size, generator=generator, dtype=dtype, device=device)[ 70 | indices 71 | ] 72 | 73 | def randint(self, low, high, size, dtype=th.long, device="cpu"): 74 | global_size, indices = self.get_global_size_and_indices(size) 75 | generator = self.get_generator(device) 76 | return th.randint( 77 | low, high, generator=generator, size=global_size, dtype=dtype, device=device 78 | )[indices] 79 | 80 | def randn_like(self, tensor): 81 | size, dtype, device = tensor.size(), tensor.dtype, tensor.device 82 | return self.randn(*size, dtype=dtype, device=device) 83 | 84 | def set_done_samples(self, done_samples): 85 | self.done_samples = done_samples 86 | self.set_seed(self.seed) 87 | 88 | def get_seed(self): 89 | return self.seed 90 | 91 | def set_seed(self, seed): 92 | self.rng_cpu.manual_seed(seed) 93 | if th.cuda.is_available(): 94 | self.rng_cuda.manual_seed(seed) 95 | 96 | 97 | class DeterministicIndividualGenerator: 98 | """ 99 | RNG to deterministically sample num_samples samples that does not depend on batch_size or mpi_machines 100 | Uses a separate rng for each sample to reduce memoery usage 101 | """ 102 | 103 | def __init__(self, num_samples, seed=0): 104 | if dist.is_initialized(): 105 | self.rank = dist.get_rank() 106 | self.world_size = dist.get_world_size() 107 | else: 108 | print("Warning: Distributed not initialised, using single rank") 109 | self.rank = 0 110 | self.world_size = 1 111 | self.num_samples = num_samples 112 | self.done_samples = 0 113 | self.seed = seed 114 | self.rng_cpu = [th.Generator() for _ in range(num_samples)] 115 | if th.cuda.is_available(): 116 | self.rng_cuda = [th.Generator(dist_util.dev()) for _ in range(num_samples)] 117 | self.set_seed(seed) 118 | 119 | def get_size_and_indices(self, size): 120 | indices = th.arange( 121 | self.done_samples + self.rank, 122 | self.done_samples + self.world_size * int(size[0]), 123 | self.world_size, 124 | ) 125 | indices = th.clamp(indices, 0, self.num_samples - 1) 126 | assert ( 127 | len(indices) == size[0] 128 | ), f"rank={self.rank}, ws={self.world_size}, l={len(indices)}, bs={size[0]}" 129 | return (1, *size[1:]), indices 130 | 131 | def get_generator(self, device): 132 | return self.rng_cpu if th.device(device).type == "cpu" else self.rng_cuda 133 | 134 | def randn(self, *size, dtype=th.float, device="cpu"): 135 | size, indices = self.get_size_and_indices(size) 136 | generator = self.get_generator(device) 137 | return th.cat( 138 | [ 139 | th.randn(*size, generator=generator[i], dtype=dtype, device=device) 140 | for i in indices 141 | ], 142 | dim=0, 143 | ) 144 | 145 | def randint(self, low, high, size, dtype=th.long, device="cpu"): 146 | size, indices = self.get_size_and_indices(size) 147 | generator = self.get_generator(device) 148 | return th.cat( 149 | [ 150 | th.randint( 151 | low, 152 | high, 153 | generator=generator[i], 154 | size=size, 155 | dtype=dtype, 156 | device=device, 157 | ) 158 | for i in indices 159 | ], 160 | dim=0, 161 | ) 162 | 163 | def randn_like(self, tensor): 164 | size, dtype, device = tensor.size(), tensor.dtype, tensor.device 165 | return self.randn(*size, dtype=dtype, device=device) 166 | 167 | def set_done_samples(self, done_samples): 168 | self.done_samples = done_samples 169 | 170 | def get_seed(self): 171 | return self.seed 172 | 173 | def set_seed(self, seed): 174 | [ 175 | rng_cpu.manual_seed(i + self.num_samples * seed) 176 | for i, rng_cpu in enumerate(self.rng_cpu) 177 | ] 178 | if th.cuda.is_available(): 179 | [ 180 | rng_cuda.manual_seed(i + self.num_samples * seed) 181 | for i, rng_cuda in enumerate(self.rng_cuda) 182 | ] 183 | -------------------------------------------------------------------------------- /iCT/cm/random_util.py: -------------------------------------------------------------------------------- 1 | import torch as th 2 | import torch.distributed as dist 3 | from . import dist_util 4 | 5 | 6 | def get_generator(generator, num_samples=0, seed=0): 7 | if generator == "dummy": 8 | return DummyGenerator() 9 | elif generator == "determ": 10 | return DeterministicGenerator(num_samples, seed) 11 | elif generator == "determ-indiv": 12 | return DeterministicIndividualGenerator(num_samples, seed) 13 | else: 14 | raise NotImplementedError 15 | 16 | 17 | class DummyGenerator: 18 | def randn(self, *args, **kwargs): 19 | return th.randn(*args, **kwargs) 20 | 21 | def randint(self, *args, **kwargs): 22 | return th.randint(*args, **kwargs) 23 | 24 | def randn_like(self, *args, **kwargs): 25 | return th.randn_like(*args, **kwargs) 26 | 27 | 28 | class DeterministicGenerator: 29 | """ 30 | RNG to deterministically sample num_samples samples that does not depend on batch_size or mpi_machines 31 | Uses a single rng and samples num_samples sized randomness and subsamples the current indices 32 | """ 33 | 34 | def __init__(self, num_samples, seed=0): 35 | if dist.is_initialized(): 36 | self.rank = dist.get_rank() 37 | self.world_size = dist.get_world_size() 38 | else: 39 | print("Warning: Distributed not initialised, using single rank") 40 | self.rank = 0 41 | self.world_size = 1 42 | self.num_samples = num_samples 43 | self.done_samples = 0 44 | self.seed = seed 45 | self.rng_cpu = th.Generator() 46 | if th.cuda.is_available(): 47 | self.rng_cuda = th.Generator(dist_util.dev()) 48 | self.set_seed(seed) 49 | 50 | def get_global_size_and_indices(self, size): 51 | global_size = (self.num_samples, *size[1:]) 52 | indices = th.arange( 53 | self.done_samples + self.rank, 54 | self.done_samples + self.world_size * int(size[0]), 55 | self.world_size, 56 | ) 57 | indices = th.clamp(indices, 0, self.num_samples - 1) 58 | assert ( 59 | len(indices) == size[0] 60 | ), f"rank={self.rank}, ws={self.world_size}, l={len(indices)}, bs={size[0]}" 61 | return global_size, indices 62 | 63 | def get_generator(self, device): 64 | return self.rng_cpu if th.device(device).type == "cpu" else self.rng_cuda 65 | 66 | def randn(self, *size, dtype=th.float, device="cpu"): 67 | global_size, indices = self.get_global_size_and_indices(size) 68 | generator = self.get_generator(device) 69 | return th.randn(*global_size, generator=generator, dtype=dtype, device=device)[ 70 | indices 71 | ] 72 | 73 | def randint(self, low, high, size, dtype=th.long, device="cpu"): 74 | global_size, indices = self.get_global_size_and_indices(size) 75 | generator = self.get_generator(device) 76 | return th.randint( 77 | low, high, generator=generator, size=global_size, dtype=dtype, device=device 78 | )[indices] 79 | 80 | def randn_like(self, tensor): 81 | size, dtype, device = tensor.size(), tensor.dtype, tensor.device 82 | return self.randn(*size, dtype=dtype, device=device) 83 | 84 | def set_done_samples(self, done_samples): 85 | self.done_samples = done_samples 86 | self.set_seed(self.seed) 87 | 88 | def get_seed(self): 89 | return self.seed 90 | 91 | def set_seed(self, seed): 92 | self.rng_cpu.manual_seed(seed) 93 | if th.cuda.is_available(): 94 | self.rng_cuda.manual_seed(seed) 95 | 96 | 97 | class DeterministicIndividualGenerator: 98 | """ 99 | RNG to deterministically sample num_samples samples that does not depend on batch_size or mpi_machines 100 | Uses a separate rng for each sample to reduce memoery usage 101 | """ 102 | 103 | def __init__(self, num_samples, seed=0): 104 | if dist.is_initialized(): 105 | self.rank = dist.get_rank() 106 | self.world_size = dist.get_world_size() 107 | else: 108 | print("Warning: Distributed not initialised, using single rank") 109 | self.rank = 0 110 | self.world_size = 1 111 | self.num_samples = num_samples 112 | self.done_samples = 0 113 | self.seed = seed 114 | self.rng_cpu = [th.Generator() for _ in range(num_samples)] 115 | if th.cuda.is_available(): 116 | self.rng_cuda = [th.Generator(dist_util.dev()) for _ in range(num_samples)] 117 | self.set_seed(seed) 118 | 119 | def get_size_and_indices(self, size): 120 | indices = th.arange( 121 | self.done_samples + self.rank, 122 | self.done_samples + self.world_size * int(size[0]), 123 | self.world_size, 124 | ) 125 | indices = th.clamp(indices, 0, self.num_samples - 1) 126 | assert ( 127 | len(indices) == size[0] 128 | ), f"rank={self.rank}, ws={self.world_size}, l={len(indices)}, bs={size[0]}" 129 | return (1, *size[1:]), indices 130 | 131 | def get_generator(self, device): 132 | return self.rng_cpu if th.device(device).type == "cpu" else self.rng_cuda 133 | 134 | def randn(self, *size, dtype=th.float, device="cpu"): 135 | size, indices = self.get_size_and_indices(size) 136 | generator = self.get_generator(device) 137 | return th.cat( 138 | [ 139 | th.randn(*size, generator=generator[i], dtype=dtype, device=device) 140 | for i in indices 141 | ], 142 | dim=0, 143 | ) 144 | 145 | def randint(self, low, high, size, dtype=th.long, device="cpu"): 146 | size, indices = self.get_size_and_indices(size) 147 | generator = self.get_generator(device) 148 | return th.cat( 149 | [ 150 | th.randint( 151 | low, 152 | high, 153 | generator=generator[i], 154 | size=size, 155 | dtype=dtype, 156 | device=device, 157 | ) 158 | for i in indices 159 | ], 160 | dim=0, 161 | ) 162 | 163 | def randn_like(self, tensor): 164 | size, dtype, device = tensor.size(), tensor.dtype, tensor.device 165 | return self.randn(*size, dtype=dtype, device=device) 166 | 167 | def set_done_samples(self, done_samples): 168 | self.done_samples = done_samples 169 | 170 | def get_seed(self): 171 | return self.seed 172 | 173 | def set_seed(self, seed): 174 | [ 175 | rng_cpu.manual_seed(i + self.num_samples * seed) 176 | for i, rng_cpu in enumerate(self.rng_cpu) 177 | ] 178 | if th.cuda.is_available(): 179 | [ 180 | rng_cuda.manual_seed(i + self.num_samples * seed) 181 | for i, rng_cuda in enumerate(self.rng_cuda) 182 | ] 183 | -------------------------------------------------------------------------------- /BCM/scripts/cm_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | import sys 5 | import copy 6 | import torch 7 | 8 | sys.path.append('../BCM/') 9 | 10 | import argparse 11 | import os 12 | 13 | from cm import dist_util, logger 14 | from cm.image_datasets import load_data 15 | from cm.resample import create_named_schedule_sampler 16 | from cm.script_util import ( 17 | model_and_diffusion_defaults, 18 | create_model_and_diffusion, 19 | cm_train_defaults, 20 | args_to_dict, 21 | add_dict_to_argparser, 22 | create_ema_and_scales_fn, 23 | ) 24 | from cm.train_util import CMTrainLoop 25 | import torch.distributed as dist 26 | import copy 27 | 28 | 29 | def main(): 30 | args = create_argparser().parse_args() 31 | 32 | dist_util.setup_dist() 33 | logger.configure(dir=os.path.join(args.save_dir, args.exp_name)) 34 | 35 | logger.log("creating model and diffusion...") 36 | ema_scale_fn = create_ema_and_scales_fn( 37 | target_ema_mode=args.target_ema_mode, 38 | start_ema=args.start_ema, 39 | scale_mode=args.scale_mode, 40 | start_scales=args.start_scales, 41 | end_scales=args.end_scales, 42 | total_steps=args.total_training_steps, 43 | distill_steps_per_iter=args.distill_steps_per_iter, 44 | ) 45 | if args.training_mode == "progdist": 46 | distillation = False 47 | elif "consistency" in args.training_mode: 48 | distillation = True 49 | else: 50 | raise ValueError(f"unknown training mode {args.training_mode}") 51 | 52 | model_and_diffusion_kwargs = args_to_dict( 53 | args, model_and_diffusion_defaults().keys() 54 | ) 55 | model_and_diffusion_kwargs["distillation"] = distillation 56 | model, diffusion = create_model_and_diffusion(**model_and_diffusion_kwargs) 57 | if args.bcf and not args.resume_checkpoint: 58 | print(f"Loading pretrained model from {args.pretrained_model_path} for BCF...") 59 | pretrained_state_dict = dist_util.load_state_dict(args.pretrained_model_path, map_location="cpu") 60 | model.load_state_dict( 61 | pretrained_state_dict, 62 | strict=False 63 | ) 64 | 65 | model.time_embed_end = copy.deepcopy(model.time_embed) # initialize t_end embedding layer from t embedding 66 | # initialize linear network to [I; 0], to make sure this initialization does not change the CM results 67 | embed_dim = model.time_embed_align.weight.data.shape[0] # Note! weight in nn.Linear is saved in its transpose 68 | model.time_embed_align.weight.data[:, : embed_dim] = torch.diag(torch.ones(embed_dim, dtype=model.time_embed_align.weight.dtype)) 69 | model.time_embed_align.weight.data[:, embed_dim:] = torch.zeros((embed_dim, embed_dim), dtype=model.time_embed_align.weight.dtype) 70 | 71 | model.to(dist_util.dev()) 72 | model.train() 73 | if args.use_fp16: 74 | model.convert_to_fp16() 75 | 76 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 77 | 78 | logger.log("creating data loader...") 79 | if args.batch_size == -1: 80 | batch_size = args.global_batch_size // dist.get_world_size() 81 | if args.global_batch_size % dist.get_world_size() != 0: 82 | logger.log( 83 | f"warning, using smaller global_batch_size of {dist.get_world_size()*batch_size} instead of {args.global_batch_size}" 84 | ) 85 | else: 86 | batch_size = args.batch_size 87 | 88 | data = load_data( 89 | data_dir=args.data_dir, 90 | batch_size=batch_size, 91 | image_size=args.image_size, 92 | class_cond=args.class_cond, 93 | ) 94 | 95 | if len(args.teacher_model_path) > 0: # path to the teacher score model. 96 | logger.log(f"loading the teacher model from {args.teacher_model_path}") 97 | teacher_model_and_diffusion_kwargs = copy.deepcopy(model_and_diffusion_kwargs) 98 | teacher_model_and_diffusion_kwargs["dropout"] = args.teacher_dropout 99 | teacher_model_and_diffusion_kwargs["distillation"] = False 100 | teacher_model, teacher_diffusion = create_model_and_diffusion( 101 | **teacher_model_and_diffusion_kwargs, 102 | ) 103 | 104 | teacher_model.load_state_dict( 105 | dist_util.load_state_dict(args.teacher_model_path, map_location="cpu"), 106 | ) 107 | 108 | teacher_model.to(dist_util.dev()) 109 | teacher_model.eval() 110 | 111 | for dst, src in zip(model.parameters(), teacher_model.parameters()): 112 | dst.data.copy_(src.data) 113 | 114 | if args.use_fp16: 115 | teacher_model.convert_to_fp16() 116 | 117 | else: 118 | teacher_model = None 119 | teacher_diffusion = None 120 | 121 | # load the target model for distillation, if path specified. 122 | 123 | logger.log("creating the target model") 124 | target_model, _ = create_model_and_diffusion( 125 | **model_and_diffusion_kwargs, 126 | ) 127 | 128 | target_model.to(dist_util.dev()) 129 | target_model.train() 130 | 131 | dist_util.sync_params(target_model.parameters()) 132 | dist_util.sync_params(target_model.buffers()) 133 | 134 | for dst, src in zip(target_model.parameters(), model.parameters()): 135 | dst.data.copy_(src.data) 136 | 137 | if args.use_fp16: 138 | target_model.convert_to_fp16() 139 | 140 | logger.log("training...") 141 | CMTrainLoop( 142 | model=model, 143 | target_model=target_model, 144 | teacher_model=teacher_model, 145 | teacher_diffusion=teacher_diffusion, 146 | training_mode=args.training_mode, 147 | ema_scale_fn=ema_scale_fn, 148 | total_training_steps=args.total_training_steps, 149 | diffusion=diffusion, 150 | data=data, 151 | batch_size=batch_size, 152 | microbatch=args.microbatch, 153 | lr=args.lr, 154 | ema_rate=args.ema_rate, 155 | log_interval=args.log_interval, 156 | save_interval=args.save_interval, 157 | resume_checkpoint=args.resume_checkpoint, 158 | use_fp16=args.use_fp16, 159 | fp16_scale_growth=args.fp16_scale_growth, 160 | schedule_sampler=schedule_sampler, 161 | weight_decay=args.weight_decay, 162 | lr_anneal_steps=args.lr_anneal_steps, 163 | ).run_loop() 164 | 165 | 166 | def create_argparser(): 167 | defaults = dict( 168 | data_dir="", 169 | schedule_sampler="uniform", 170 | lr=1e-4, 171 | weight_decay=0.0, 172 | lr_anneal_steps=0, 173 | global_batch_size=4096, 174 | batch_size=-1, 175 | microbatch=-1, # -1 disables microbatches 176 | ema_rate="0.9999", # comma-separated list of EMA values 177 | log_interval=10, 178 | save_interval=10000, 179 | save_dir='./checkpoints', 180 | exp_name='ict', 181 | resume_checkpoint="", 182 | use_fp16=False, 183 | fp16_scale_growth=1e-3, 184 | bcf=False, 185 | pretrained_model_path="", 186 | 187 | ) 188 | defaults.update(model_and_diffusion_defaults()) 189 | defaults.update(cm_train_defaults()) 190 | parser = argparse.ArgumentParser() 191 | add_dict_to_argparser(parser, defaults) 192 | return parser 193 | 194 | 195 | if __name__ == "__main__": 196 | main() 197 | -------------------------------------------------------------------------------- /BCM/cm/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | from scipy.stats import norm 6 | import torch.distributed as dist 7 | 8 | 9 | def create_named_schedule_sampler(name, diffusion): 10 | """ 11 | Create a ScheduleSampler from a library of pre-defined samplers. 12 | 13 | :param name: the name of the sampler. 14 | :param diffusion: the diffusion object to sample for. 15 | """ 16 | if name == "uniform": 17 | return UniformSampler(diffusion) 18 | elif name == "loss-second-moment": 19 | return LossSecondMomentResampler(diffusion) 20 | elif name == "lognormal": 21 | return LogNormalSampler() 22 | else: 23 | raise NotImplementedError(f"unknown schedule sampler: {name}") 24 | 25 | 26 | class ScheduleSampler(ABC): 27 | """ 28 | A distribution over timesteps in the diffusion process, intended to reduce 29 | variance of the objective. 30 | 31 | By default, samplers perform unbiased importance sampling, in which the 32 | objective's mean is unchanged. 33 | However, subclasses may override sample() to change how the resampled 34 | terms are reweighted, allowing for actual changes in the objective. 35 | """ 36 | 37 | @abstractmethod 38 | def weights(self): 39 | """ 40 | Get a numpy array of weights, one per diffusion step. 41 | 42 | The weights needn't be normalized, but must be positive. 43 | """ 44 | 45 | def sample(self, batch_size, device): 46 | """ 47 | Importance-sample timesteps for a batch. 48 | 49 | :param batch_size: the number of timesteps. 50 | :param device: the torch device to save to. 51 | :return: a tuple (timesteps, weights): 52 | - timesteps: a tensor of timestep indices. 53 | - weights: a tensor of weights to scale the resulting losses. 54 | """ 55 | w = self.weights() 56 | p = w / np.sum(w) 57 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 58 | indices = th.from_numpy(indices_np).long().to(device) 59 | weights_np = 1 / (len(p) * p[indices_np]) 60 | weights = th.from_numpy(weights_np).float().to(device) 61 | return indices, weights 62 | 63 | 64 | class UniformSampler(ScheduleSampler): 65 | def __init__(self, diffusion): 66 | self.diffusion = diffusion 67 | self._weights = np.ones([diffusion.num_timesteps]) 68 | 69 | def weights(self): 70 | return self._weights 71 | 72 | 73 | class LossAwareSampler(ScheduleSampler): 74 | def update_with_local_losses(self, local_ts, local_losses): 75 | """ 76 | Update the reweighting using losses from a model. 77 | 78 | Call this method from each rank with a batch of timesteps and the 79 | corresponding losses for each of those timesteps. 80 | This method will perform synchronization to make sure all of the ranks 81 | maintain the exact same reweighting. 82 | 83 | :param local_ts: an integer Tensor of timesteps. 84 | :param local_losses: a 1D Tensor of losses. 85 | """ 86 | batch_sizes = [ 87 | th.tensor([0], dtype=th.int32, device=local_ts.device) 88 | for _ in range(dist.get_world_size()) 89 | ] 90 | dist.all_gather( 91 | batch_sizes, 92 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 93 | ) 94 | 95 | # Pad all_gather batches to be the maximum batch size. 96 | batch_sizes = [x.item() for x in batch_sizes] 97 | max_bs = max(batch_sizes) 98 | 99 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 100 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 101 | dist.all_gather(timestep_batches, local_ts) 102 | dist.all_gather(loss_batches, local_losses) 103 | timesteps = [ 104 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 105 | ] 106 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 107 | self.update_with_all_losses(timesteps, losses) 108 | 109 | @abstractmethod 110 | def update_with_all_losses(self, ts, losses): 111 | """ 112 | Update the reweighting using losses from a model. 113 | 114 | Sub-classes should override this method to update the reweighting 115 | using losses from the model. 116 | 117 | This method directly updates the reweighting without synchronizing 118 | between workers. It is called by update_with_local_losses from all 119 | ranks with identical arguments. Thus, it should have deterministic 120 | behavior to maintain state across workers. 121 | 122 | :param ts: a list of int timesteps. 123 | :param losses: a list of float losses, one per timestep. 124 | """ 125 | 126 | 127 | class LossSecondMomentResampler(LossAwareSampler): 128 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 129 | self.diffusion = diffusion 130 | self.history_per_term = history_per_term 131 | self.uniform_prob = uniform_prob 132 | self._loss_history = np.zeros( 133 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 134 | ) 135 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 136 | 137 | def weights(self): 138 | if not self._warmed_up(): 139 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 140 | weights = np.sqrt(np.mean(self._loss_history**2, axis=-1)) 141 | weights /= np.sum(weights) 142 | weights *= 1 - self.uniform_prob 143 | weights += self.uniform_prob / len(weights) 144 | return weights 145 | 146 | def update_with_all_losses(self, ts, losses): 147 | for t, loss in zip(ts, losses): 148 | if self._loss_counts[t] == self.history_per_term: 149 | # Shift out the oldest loss term. 150 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 151 | self._loss_history[t, -1] = loss 152 | else: 153 | self._loss_history[t, self._loss_counts[t]] = loss 154 | self._loss_counts[t] += 1 155 | 156 | def _warmed_up(self): 157 | return (self._loss_counts == self.history_per_term).all() 158 | 159 | 160 | class LogNormalSampler: 161 | def __init__(self, p_mean=-1.1, p_std=2.0, even=False): 162 | self.p_mean = p_mean 163 | self.p_std = p_std 164 | self.even = even 165 | if self.even: 166 | self.inv_cdf = lambda x: norm.ppf(x, loc=p_mean, scale=p_std) 167 | self.rank, self.size = dist.get_rank(), dist.get_world_size() 168 | 169 | def sample(self, bs, device): 170 | if self.even: 171 | # buckets = [1/G] 172 | start_i, end_i = self.rank * bs, (self.rank + 1) * bs 173 | global_batch_size = self.size * bs 174 | locs = (th.arange(start_i, end_i) + th.rand(bs)) / global_batch_size 175 | log_sigmas = th.tensor(self.inv_cdf(locs), dtype=th.float32, device=device) 176 | else: 177 | log_sigmas = self.p_mean + self.p_std * th.randn(bs, device=device) 178 | sigmas = th.exp(log_sigmas) 179 | weights = th.ones_like(sigmas) 180 | return sigmas, weights 181 | -------------------------------------------------------------------------------- /iCT/cm/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | from scipy.stats import norm 6 | import torch.distributed as dist 7 | 8 | 9 | def create_named_schedule_sampler(name, diffusion): 10 | """ 11 | Create a ScheduleSampler from a library of pre-defined samplers. 12 | 13 | :param name: the name of the sampler. 14 | :param diffusion: the diffusion object to sample for. 15 | """ 16 | if name == "uniform": 17 | return UniformSampler(diffusion) 18 | elif name == "loss-second-moment": 19 | return LossSecondMomentResampler(diffusion) 20 | elif name == "lognormal": 21 | return LogNormalSampler() 22 | else: 23 | raise NotImplementedError(f"unknown schedule sampler: {name}") 24 | 25 | 26 | class ScheduleSampler(ABC): 27 | """ 28 | A distribution over timesteps in the diffusion process, intended to reduce 29 | variance of the objective. 30 | 31 | By default, samplers perform unbiased importance sampling, in which the 32 | objective's mean is unchanged. 33 | However, subclasses may override sample() to change how the resampled 34 | terms are reweighted, allowing for actual changes in the objective. 35 | """ 36 | 37 | @abstractmethod 38 | def weights(self): 39 | """ 40 | Get a numpy array of weights, one per diffusion step. 41 | 42 | The weights needn't be normalized, but must be positive. 43 | """ 44 | 45 | def sample(self, batch_size, device): 46 | """ 47 | Importance-sample timesteps for a batch. 48 | 49 | :param batch_size: the number of timesteps. 50 | :param device: the torch device to save to. 51 | :return: a tuple (timesteps, weights): 52 | - timesteps: a tensor of timestep indices. 53 | - weights: a tensor of weights to scale the resulting losses. 54 | """ 55 | w = self.weights() 56 | p = w / np.sum(w) 57 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 58 | indices = th.from_numpy(indices_np).long().to(device) 59 | weights_np = 1 / (len(p) * p[indices_np]) 60 | weights = th.from_numpy(weights_np).float().to(device) 61 | return indices, weights 62 | 63 | 64 | class UniformSampler(ScheduleSampler): 65 | def __init__(self, diffusion): 66 | self.diffusion = diffusion 67 | self._weights = np.ones([diffusion.num_timesteps]) 68 | 69 | def weights(self): 70 | return self._weights 71 | 72 | 73 | class LossAwareSampler(ScheduleSampler): 74 | def update_with_local_losses(self, local_ts, local_losses): 75 | """ 76 | Update the reweighting using losses from a model. 77 | 78 | Call this method from each rank with a batch of timesteps and the 79 | corresponding losses for each of those timesteps. 80 | This method will perform synchronization to make sure all of the ranks 81 | maintain the exact same reweighting. 82 | 83 | :param local_ts: an integer Tensor of timesteps. 84 | :param local_losses: a 1D Tensor of losses. 85 | """ 86 | batch_sizes = [ 87 | th.tensor([0], dtype=th.int32, device=local_ts.device) 88 | for _ in range(dist.get_world_size()) 89 | ] 90 | dist.all_gather( 91 | batch_sizes, 92 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 93 | ) 94 | 95 | # Pad all_gather batches to be the maximum batch size. 96 | batch_sizes = [x.item() for x in batch_sizes] 97 | max_bs = max(batch_sizes) 98 | 99 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 100 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 101 | dist.all_gather(timestep_batches, local_ts) 102 | dist.all_gather(loss_batches, local_losses) 103 | timesteps = [ 104 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 105 | ] 106 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 107 | self.update_with_all_losses(timesteps, losses) 108 | 109 | @abstractmethod 110 | def update_with_all_losses(self, ts, losses): 111 | """ 112 | Update the reweighting using losses from a model. 113 | 114 | Sub-classes should override this method to update the reweighting 115 | using losses from the model. 116 | 117 | This method directly updates the reweighting without synchronizing 118 | between workers. It is called by update_with_local_losses from all 119 | ranks with identical arguments. Thus, it should have deterministic 120 | behavior to maintain state across workers. 121 | 122 | :param ts: a list of int timesteps. 123 | :param losses: a list of float losses, one per timestep. 124 | """ 125 | 126 | 127 | class LossSecondMomentResampler(LossAwareSampler): 128 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 129 | self.diffusion = diffusion 130 | self.history_per_term = history_per_term 131 | self.uniform_prob = uniform_prob 132 | self._loss_history = np.zeros( 133 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 134 | ) 135 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 136 | 137 | def weights(self): 138 | if not self._warmed_up(): 139 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 140 | weights = np.sqrt(np.mean(self._loss_history**2, axis=-1)) 141 | weights /= np.sum(weights) 142 | weights *= 1 - self.uniform_prob 143 | weights += self.uniform_prob / len(weights) 144 | return weights 145 | 146 | def update_with_all_losses(self, ts, losses): 147 | for t, loss in zip(ts, losses): 148 | if self._loss_counts[t] == self.history_per_term: 149 | # Shift out the oldest loss term. 150 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 151 | self._loss_history[t, -1] = loss 152 | else: 153 | self._loss_history[t, self._loss_counts[t]] = loss 154 | self._loss_counts[t] += 1 155 | 156 | def _warmed_up(self): 157 | return (self._loss_counts == self.history_per_term).all() 158 | 159 | 160 | class LogNormalSampler: 161 | def __init__(self, p_mean=-1.1, p_std=2.0, even=False): 162 | self.p_mean = p_mean 163 | self.p_std = p_std 164 | self.even = even 165 | if self.even: 166 | self.inv_cdf = lambda x: norm.ppf(x, loc=p_mean, scale=p_std) 167 | self.rank, self.size = dist.get_rank(), dist.get_world_size() 168 | 169 | def sample(self, bs, device): 170 | if self.even: 171 | # buckets = [1/G] 172 | start_i, end_i = self.rank * bs, (self.rank + 1) * bs 173 | global_batch_size = self.size * bs 174 | locs = (th.arange(start_i, end_i) + th.rand(bs)) / global_batch_size 175 | log_sigmas = th.tensor(self.inv_cdf(locs), dtype=th.float32, device=device) 176 | else: 177 | log_sigmas = self.p_mean + self.p_std * th.randn(bs, device=device) 178 | sigmas = th.exp(log_sigmas) 179 | weights = th.ones_like(sigmas) 180 | return sigmas, weights 181 | -------------------------------------------------------------------------------- /BCM/cm/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 | val=False 21 | ): 22 | """ 23 | For a dataset, create a generator over (images, kwargs) pairs. 24 | 25 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 26 | more keys, each of which map to a batched Tensor of their own. 27 | The kwargs dict can be used for class labels, in which case the key is "y" 28 | and the values are integer tensors of class labels. 29 | 30 | :param data_dir: a dataset directory. 31 | :param batch_size: the batch size of each returned pair. 32 | :param image_size: the size to which images are resized. 33 | :param class_cond: if True, include a "y" key in returned dicts for class 34 | label. If classes are not available and this is true, an 35 | exception will be raised. 36 | :param deterministic: if True, yield results in a deterministic order. 37 | :param random_crop: if True, randomly crop the images for augmentation. 38 | :param random_flip: if True, randomly flip the images for augmentation. 39 | """ 40 | if not data_dir: 41 | raise ValueError("unspecified data directory") 42 | classes = None 43 | all_files = _list_image_files_recursively(data_dir) 44 | if class_cond: 45 | if not val: # imagenet的training set可以直接从图片名字读出label 46 | # Assume classes are the first part of the filename, 47 | # before an underscore. 48 | class_names = [bf.basename(path).split("_")[0] for path in all_files] 49 | sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 50 | classes = [sorted_classes[x] for x in class_names] 51 | else: 52 | val_label_text_path = 'datasets/imagenet_val_label.txt' 53 | val_label_dict = {} 54 | with open(val_label_text_path, "r") as f: 55 | for line in f.readlines(): 56 | line = line.strip('\n') 57 | img, label = line.split(' ') 58 | val_label_dict[img] = int(label) 59 | classes = [val_label_dict[path.split('/')[-1]] for path in all_files] 60 | 61 | dataset = ImageDataset( 62 | image_size, 63 | all_files, 64 | classes=classes, 65 | shard=MPI.COMM_WORLD.Get_rank(), 66 | num_shards=MPI.COMM_WORLD.Get_size(), 67 | random_crop=random_crop, 68 | random_flip=random_flip, 69 | ) 70 | if deterministic: 71 | loader = DataLoader( 72 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 73 | ) 74 | else: 75 | loader = DataLoader( 76 | dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True 77 | ) 78 | while True: 79 | yield from loader 80 | 81 | 82 | def _list_image_files_recursively(data_dir): 83 | results = [] 84 | for entry in sorted(bf.listdir(data_dir)): 85 | full_path = bf.join(data_dir, entry) 86 | ext = entry.split(".")[-1] 87 | if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]: 88 | results.append(full_path) 89 | elif bf.isdir(full_path): 90 | results.extend(_list_image_files_recursively(full_path)) 91 | return results 92 | 93 | 94 | class ImageDataset(Dataset): 95 | def __init__( 96 | self, 97 | resolution, 98 | image_paths, 99 | classes=None, 100 | shard=0, 101 | num_shards=1, 102 | random_crop=False, 103 | random_flip=True, 104 | ): 105 | super().__init__() 106 | self.resolution = resolution 107 | self.local_images = image_paths[shard:][::num_shards] 108 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 109 | self.random_crop = random_crop 110 | self.random_flip = random_flip 111 | self.save_jpg = False 112 | self.jpg_quality = 20 113 | self.jpg_save_path = f'/mnt/petrelfs/liliangchen/try/consistency_model_torch/bcm/playground/imagenet64_val_jpg_qf{self.jpg_quality}' 114 | 115 | def __len__(self): 116 | return len(self.local_images) 117 | 118 | def __getitem__(self, idx): 119 | path = self.local_images[idx] 120 | with bf.BlobFile(path, "rb") as f: 121 | pil_image = Image.open(f) 122 | pil_image.load() 123 | pil_image = pil_image.convert("RGB") 124 | 125 | if self.random_crop: 126 | arr = random_crop_arr(pil_image, self.resolution) 127 | else: 128 | arr = center_crop_arr(pil_image, self.resolution) 129 | 130 | if self.save_jpg: 131 | pic_name = path.split('/')[-1] 132 | temp = Image.fromarray(arr) 133 | temp.save(f'{self.jpg_save_path}/{pic_name}', quality=self.jpg_quality) 134 | 135 | if self.random_flip and random.random() < 0.5: 136 | arr = arr[:, ::-1] 137 | 138 | arr = arr.astype(np.float32) / 127.5 - 1 139 | 140 | out_dict = {} 141 | if self.local_classes is not None: 142 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 143 | return np.transpose(arr, [2, 0, 1]), out_dict 144 | 145 | 146 | def center_crop_arr(pil_image, image_size): 147 | # We are not on a new enough PIL to support the `reducing_gap` 148 | # argument, which uses BOX downsampling at powers of two first. 149 | # Thus, we do it by hand to improve downsample quality. 150 | while min(*pil_image.size) >= 2 * image_size: 151 | pil_image = pil_image.resize( 152 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 153 | ) 154 | 155 | scale = image_size / min(*pil_image.size) 156 | pil_image = pil_image.resize( 157 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 158 | ) 159 | 160 | arr = np.array(pil_image) 161 | crop_y = (arr.shape[0] - image_size) // 2 162 | crop_x = (arr.shape[1] - image_size) // 2 163 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 164 | 165 | 166 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): 167 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 168 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 169 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 170 | 171 | # We are not on a new enough PIL to support the `reducing_gap` 172 | # argument, which uses BOX downsampling at powers of two first. 173 | # Thus, we do it by hand to improve downsample quality. 174 | while min(*pil_image.size) >= 2 * smaller_dim_size: 175 | pil_image = pil_image.resize( 176 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 177 | ) 178 | 179 | scale = smaller_dim_size / min(*pil_image.size) 180 | pil_image = pil_image.resize( 181 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 182 | ) 183 | 184 | arr = np.array(pil_image) 185 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 186 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 187 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 188 | -------------------------------------------------------------------------------- /BCM/cm/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from .karras_diffusion import KarrasDenoiser 4 | from .unet_bcf import UNetModel 5 | import numpy as np 6 | 7 | NUM_CLASSES = 1000 8 | 9 | 10 | def cm_train_defaults(): 11 | return dict( 12 | teacher_model_path="", 13 | teacher_dropout=0.1, 14 | training_mode="consistency_distillation", 15 | target_ema_mode="fixed", 16 | scale_mode="fixed", 17 | total_training_steps=600000, 18 | start_ema=0.0, 19 | start_scales=40, 20 | end_scales=40, 21 | distill_steps_per_iter=50000, 22 | loss_norm="lpips", 23 | ) 24 | 25 | 26 | def model_and_diffusion_defaults(): 27 | """ 28 | Defaults for image training. 29 | """ 30 | res = dict( 31 | sigma_min=0.002, 32 | sigma_max=80.0, 33 | image_size=64, 34 | num_channels=128, 35 | num_res_blocks=2, 36 | num_heads=4, 37 | num_heads_upsample=-1, 38 | num_head_channels=-1, 39 | attention_resolutions="32,16,8", 40 | channel_mult="", 41 | dropout=0.0, 42 | class_cond=False, 43 | use_checkpoint=False, 44 | use_scale_shift_norm=True, 45 | resblock_updown=False, 46 | use_fp16=False, 47 | use_new_attention_order=False, 48 | learn_sigma=False, 49 | weight_schedule="karras", 50 | ) 51 | return res 52 | 53 | 54 | def create_model_and_diffusion( 55 | image_size, 56 | class_cond, 57 | learn_sigma, 58 | num_channels, 59 | num_res_blocks, 60 | channel_mult, 61 | num_heads, 62 | num_head_channels, 63 | num_heads_upsample, 64 | attention_resolutions, 65 | dropout, 66 | use_checkpoint, 67 | use_scale_shift_norm, 68 | resblock_updown, 69 | use_fp16, 70 | use_new_attention_order, 71 | weight_schedule, 72 | sigma_min=0.002, 73 | sigma_max=80.0, 74 | distillation=False, 75 | ): 76 | model = create_model( 77 | image_size, 78 | num_channels, 79 | num_res_blocks, 80 | channel_mult=channel_mult, 81 | learn_sigma=learn_sigma, 82 | class_cond=class_cond, 83 | use_checkpoint=use_checkpoint, 84 | attention_resolutions=attention_resolutions, 85 | num_heads=num_heads, 86 | num_head_channels=num_head_channels, 87 | num_heads_upsample=num_heads_upsample, 88 | use_scale_shift_norm=use_scale_shift_norm, 89 | dropout=dropout, 90 | resblock_updown=resblock_updown, 91 | use_fp16=use_fp16, 92 | use_new_attention_order=use_new_attention_order, 93 | ) 94 | diffusion = KarrasDenoiser( 95 | sigma_data=0.5, 96 | sigma_max=sigma_max, 97 | sigma_min=sigma_min, 98 | distillation=distillation, 99 | weight_schedule=weight_schedule, 100 | ) 101 | return model, diffusion 102 | 103 | 104 | def create_model( 105 | image_size, 106 | num_channels, 107 | num_res_blocks, 108 | channel_mult="", 109 | learn_sigma=False, 110 | class_cond=False, 111 | use_checkpoint=False, 112 | attention_resolutions="16", 113 | num_heads=1, 114 | num_head_channels=-1, 115 | num_heads_upsample=-1, 116 | use_scale_shift_norm=False, 117 | dropout=0, 118 | resblock_updown=False, 119 | use_fp16=False, 120 | use_new_attention_order=False, 121 | ): 122 | if channel_mult == "": 123 | if image_size == 512: 124 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 125 | elif image_size == 256: 126 | channel_mult = (1, 1, 2, 2, 4, 4) 127 | elif image_size == 128: 128 | channel_mult = (1, 1, 2, 3, 4) 129 | elif image_size == 64: 130 | channel_mult = (1, 2, 3, 4) 131 | else: 132 | raise ValueError(f"unsupported image size: {image_size}") 133 | else: 134 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 135 | 136 | attention_ds = [] 137 | for res in attention_resolutions.split(","): 138 | attention_ds.append(image_size // int(res)) 139 | 140 | return UNetModel( 141 | image_size=image_size, 142 | in_channels=3, 143 | model_channels=num_channels, 144 | out_channels=(3 if not learn_sigma else 6), 145 | num_res_blocks=num_res_blocks, 146 | attention_resolutions=tuple(attention_ds), 147 | dropout=dropout, 148 | channel_mult=channel_mult, 149 | num_classes=(NUM_CLASSES if class_cond else None), 150 | use_checkpoint=use_checkpoint, 151 | use_fp16=use_fp16, 152 | num_heads=num_heads, 153 | num_head_channels=num_head_channels, 154 | num_heads_upsample=num_heads_upsample, 155 | use_scale_shift_norm=use_scale_shift_norm, 156 | resblock_updown=resblock_updown, 157 | use_new_attention_order=use_new_attention_order, 158 | ) 159 | 160 | 161 | def create_ema_and_scales_fn( 162 | target_ema_mode, 163 | start_ema, 164 | scale_mode, 165 | start_scales, 166 | end_scales, 167 | total_steps, 168 | distill_steps_per_iter, 169 | ): 170 | def ema_and_scales_fn(step): 171 | if target_ema_mode == "fixed" and scale_mode == "fixed": 172 | target_ema = start_ema 173 | scales = start_scales 174 | elif target_ema_mode == "fixed" and scale_mode == "progressive": 175 | target_ema = start_ema 176 | scales = np.ceil( 177 | np.sqrt( 178 | (step / total_steps) * ((end_scales + 1) ** 2 - start_scales ** 2) 179 | + start_scales ** 2 180 | ) 181 | - 1 182 | ).astype(np.int32) 183 | scales = np.maximum(scales, 1) 184 | scales = scales + 1 185 | 186 | elif target_ema_mode == "adaptive" and scale_mode == "progressive": 187 | 188 | Kp = np.floor(total_steps / (np.log2(end_scales / start_scales) + 1)) 189 | scales = np.minimum(end_scales, start_scales * 2 ** np.floor(step / Kp)).astype(np.int32) + 1 190 | target_ema = 0. 191 | 192 | elif target_ema_mode == "fixed" and scale_mode == "progdist": 193 | distill_stage = step // distill_steps_per_iter 194 | scales = start_scales // (2 ** distill_stage) 195 | scales = np.maximum(scales, 2) 196 | 197 | sub_stage = np.maximum( 198 | step - distill_steps_per_iter * (np.log2(start_scales) - 1), 199 | 0, 200 | ) 201 | sub_stage = sub_stage // (distill_steps_per_iter * 2) 202 | sub_scales = 2 // (2 ** sub_stage) 203 | sub_scales = np.maximum(sub_scales, 1) 204 | 205 | scales = np.where(scales == 2, sub_scales, scales) 206 | 207 | target_ema = 1.0 208 | else: 209 | raise NotImplementedError 210 | 211 | return float(target_ema), int(scales) 212 | 213 | return ema_and_scales_fn 214 | 215 | 216 | def add_dict_to_argparser(parser, default_dict): 217 | for k, v in default_dict.items(): 218 | v_type = type(v) 219 | if v is None: 220 | v_type = str 221 | elif isinstance(v, bool): 222 | v_type = str2bool 223 | parser.add_argument(f"--{k}", default=v, type=v_type) 224 | 225 | 226 | def args_to_dict(args, keys): 227 | return {k: getattr(args, k) for k in keys} 228 | 229 | 230 | def str2bool(v): 231 | """ 232 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 233 | """ 234 | if isinstance(v, bool): 235 | return v 236 | if v.lower() in ("yes", "true", "t", "y", "1"): 237 | return True 238 | elif v.lower() in ("no", "false", "f", "n", "0"): 239 | return False 240 | else: 241 | raise argparse.ArgumentTypeError("boolean value expected") 242 | -------------------------------------------------------------------------------- /iCT/cm/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from .karras_diffusion import KarrasDenoiser 4 | from .unet import UNetModel 5 | import numpy as np 6 | 7 | NUM_CLASSES = 1000 8 | 9 | 10 | def cm_train_defaults(): 11 | return dict( 12 | teacher_model_path="", 13 | teacher_dropout=0.1, 14 | training_mode="consistency_distillation", 15 | target_ema_mode="fixed", 16 | scale_mode="fixed", 17 | total_training_steps=600000, 18 | start_ema=0.0, 19 | start_scales=40, 20 | end_scales=40, 21 | distill_steps_per_iter=50000, 22 | loss_norm="lpips", 23 | ) 24 | 25 | 26 | def model_and_diffusion_defaults(): 27 | """ 28 | Defaults for image training. 29 | """ 30 | res = dict( 31 | sigma_min=0.002, 32 | sigma_max=80.0, 33 | image_size=64, 34 | num_channels=128, 35 | num_res_blocks=2, 36 | num_heads=4, 37 | num_heads_upsample=-1, 38 | num_head_channels=-1, 39 | attention_resolutions="32,16,8", 40 | channel_mult="", 41 | dropout=0.0, 42 | class_cond=False, 43 | use_checkpoint=False, 44 | use_scale_shift_norm=True, 45 | resblock_updown=False, 46 | use_fp16=False, 47 | use_new_attention_order=False, 48 | learn_sigma=False, 49 | weight_schedule="karras", 50 | ) 51 | return res 52 | 53 | 54 | def create_model_and_diffusion( 55 | image_size, 56 | class_cond, 57 | learn_sigma, 58 | num_channels, 59 | num_res_blocks, 60 | channel_mult, 61 | num_heads, 62 | num_head_channels, 63 | num_heads_upsample, 64 | attention_resolutions, 65 | dropout, 66 | use_checkpoint, 67 | use_scale_shift_norm, 68 | resblock_updown, 69 | use_fp16, 70 | use_new_attention_order, 71 | weight_schedule, 72 | sigma_min=0.002, 73 | sigma_max=80.0, 74 | distillation=False, 75 | ): 76 | model = create_model( 77 | image_size, 78 | num_channels, 79 | num_res_blocks, 80 | channel_mult=channel_mult, 81 | learn_sigma=learn_sigma, 82 | class_cond=class_cond, 83 | use_checkpoint=use_checkpoint, 84 | attention_resolutions=attention_resolutions, 85 | num_heads=num_heads, 86 | num_head_channels=num_head_channels, 87 | num_heads_upsample=num_heads_upsample, 88 | use_scale_shift_norm=use_scale_shift_norm, 89 | dropout=dropout, 90 | resblock_updown=resblock_updown, 91 | use_fp16=use_fp16, 92 | use_new_attention_order=use_new_attention_order, 93 | ) 94 | diffusion = KarrasDenoiser( 95 | sigma_data=0.5, 96 | sigma_max=sigma_max, 97 | sigma_min=sigma_min, 98 | distillation=distillation, 99 | weight_schedule=weight_schedule, 100 | ) 101 | return model, diffusion 102 | 103 | 104 | def create_model( 105 | image_size, 106 | num_channels, 107 | num_res_blocks, 108 | channel_mult="", 109 | learn_sigma=False, 110 | class_cond=False, 111 | use_checkpoint=False, 112 | attention_resolutions="16", 113 | num_heads=1, 114 | num_head_channels=-1, 115 | num_heads_upsample=-1, 116 | use_scale_shift_norm=False, 117 | dropout=0, 118 | resblock_updown=False, 119 | use_fp16=False, 120 | use_new_attention_order=False, 121 | ): 122 | if channel_mult == "": 123 | if image_size == 512: 124 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 125 | elif image_size == 256: 126 | channel_mult = (1, 1, 2, 2, 4, 4) 127 | elif image_size == 128: 128 | channel_mult = (1, 1, 2, 3, 4) 129 | elif image_size == 64: 130 | channel_mult = (1, 2, 3, 4) 131 | else: 132 | raise ValueError(f"unsupported image size: {image_size}") 133 | else: 134 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 135 | 136 | attention_ds = [] 137 | for res in attention_resolutions.split(","): 138 | attention_ds.append(image_size // int(res)) 139 | 140 | return UNetModel( 141 | image_size=image_size, 142 | in_channels=3, 143 | model_channels=num_channels, 144 | out_channels=(3 if not learn_sigma else 6), 145 | num_res_blocks=num_res_blocks, 146 | attention_resolutions=tuple(attention_ds), 147 | dropout=dropout, 148 | channel_mult=channel_mult, 149 | num_classes=(NUM_CLASSES if class_cond else None), 150 | use_checkpoint=use_checkpoint, 151 | use_fp16=use_fp16, 152 | num_heads=num_heads, 153 | num_head_channels=num_head_channels, 154 | num_heads_upsample=num_heads_upsample, 155 | use_scale_shift_norm=use_scale_shift_norm, 156 | resblock_updown=resblock_updown, 157 | use_new_attention_order=use_new_attention_order, 158 | ) 159 | 160 | 161 | def create_ema_and_scales_fn( 162 | target_ema_mode, 163 | start_ema, 164 | scale_mode, 165 | start_scales, 166 | end_scales, 167 | total_steps, 168 | distill_steps_per_iter, 169 | ): 170 | def ema_and_scales_fn(step): 171 | if target_ema_mode == "fixed" and scale_mode == "fixed": 172 | target_ema = start_ema 173 | scales = start_scales 174 | elif target_ema_mode == "fixed" and scale_mode == "progressive": 175 | target_ema = start_ema 176 | scales = np.ceil( 177 | np.sqrt( 178 | (step / total_steps) * ((end_scales + 1) ** 2 - start_scales ** 2) 179 | + start_scales ** 2 180 | ) 181 | - 1 182 | ).astype(np.int32) 183 | scales = np.maximum(scales, 1) 184 | scales = scales + 1 185 | 186 | elif target_ema_mode == "adaptive" and scale_mode == "v1": 187 | scales = np.ceil( 188 | np.sqrt( 189 | (step / total_steps) * ((end_scales + 1) ** 2 - start_scales**2) 190 | + start_scales**2 191 | ) 192 | - 1 193 | ).astype(np.int32) 194 | scales = np.maximum(scales, 1) 195 | scales = scales + 1 196 | target_ema = 0. 197 | 198 | elif target_ema_mode == "adaptive" and scale_mode == "progressive": 199 | Kp = np.floor(total_steps / (np.log2(end_scales / start_scales) + 1)) 200 | scales = np.minimum(end_scales, start_scales * 2 ** np.floor(step / Kp)).astype(np.int32) + 1 201 | target_ema = 0. 202 | 203 | 204 | elif target_ema_mode == 'adaptive' and scale_mode == "exponential": 205 | rho = -1.0 206 | alpha1 = 1 / (start_scales + 1) 207 | alpha2 = 1 / (end_scales + 1) 208 | scales = alpha1 ** rho + step / total_steps * (alpha2 ** rho - alpha1 ** rho) 209 | scales = (np.floor(1 / scales ** (1 / rho)) + 1).astype(np.int32) 210 | target_ema = 0. 211 | 212 | elif target_ema_mode == "fixed" and scale_mode == "progdist": 213 | distill_stage = step // distill_steps_per_iter 214 | scales = start_scales // (2 ** distill_stage) 215 | scales = np.maximum(scales, 2) 216 | 217 | sub_stage = np.maximum( 218 | step - distill_steps_per_iter * (np.log2(start_scales) - 1), 219 | 0, 220 | ) 221 | sub_stage = sub_stage // (distill_steps_per_iter * 2) 222 | sub_scales = 2 // (2 ** sub_stage) 223 | sub_scales = np.maximum(sub_scales, 1) 224 | 225 | scales = np.where(scales == 2, sub_scales, scales) 226 | 227 | target_ema = 1.0 228 | else: 229 | raise NotImplementedError 230 | 231 | return float(target_ema), int(scales) 232 | 233 | return ema_and_scales_fn 234 | 235 | 236 | def add_dict_to_argparser(parser, default_dict): 237 | for k, v in default_dict.items(): 238 | v_type = type(v) 239 | if v is None: 240 | v_type = str 241 | elif isinstance(v, bool): 242 | v_type = str2bool 243 | parser.add_argument(f"--{k}", default=v, type=v_type) 244 | 245 | 246 | def args_to_dict(args, keys): 247 | return {k: getattr(args, k) for k in keys} 248 | 249 | 250 | def str2bool(v): 251 | """ 252 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 253 | """ 254 | if isinstance(v, bool): 255 | return v 256 | if v.lower() in ("yes", "true", "t", "y", "1"): 257 | return True 258 | elif v.lower() in ("no", "false", "f", "n", "0"): 259 | return False 260 | else: 261 | raise argparse.ArgumentTypeError("boolean value expected") 262 | -------------------------------------------------------------------------------- /iCT/cm/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 | if True: 100 | state_dict = model.state_dict() 101 | for master_param, (param_group, _) in zip( 102 | master_params, param_groups_and_shapes 103 | ): 104 | for (name, _), unflat_master_param in zip( 105 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 106 | ): 107 | assert name in state_dict 108 | state_dict[name] = unflat_master_param 109 | else: 110 | state_dict = model.state_dict() 111 | for i, (name, _value) in enumerate(model.named_parameters()): 112 | assert name in state_dict 113 | state_dict[name] = master_params[i] 114 | return state_dict 115 | 116 | 117 | def state_dict_to_master_params(model, state_dict, use_fp16): 118 | # if use_fp16: 119 | if True: 120 | named_model_params = [ 121 | (name, state_dict[name]) for name, _ in model.named_parameters() 122 | ] 123 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 124 | master_params = make_master_params(param_groups_and_shapes) 125 | else: 126 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 127 | return master_params 128 | 129 | 130 | def zero_master_grads(master_params): 131 | for param in master_params: 132 | param.grad = None 133 | 134 | 135 | def zero_grad(model_params): 136 | for param in model_params: 137 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 138 | if param.grad is not None: 139 | param.grad.detach_() 140 | param.grad.zero_() 141 | 142 | 143 | def param_grad_or_zeros(param): 144 | if param.grad is not None: 145 | return param.grad.data.detach() 146 | else: 147 | return th.zeros_like(param) 148 | 149 | 150 | class MixedPrecisionTrainer: 151 | def __init__( 152 | self, 153 | *, 154 | model, 155 | use_fp16=False, 156 | fp16_scale_growth=1e-3, 157 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 158 | ): 159 | self.model = model 160 | self.use_fp16 = use_fp16 161 | self.fp16_scale_growth = fp16_scale_growth 162 | 163 | self.model_params = list(self.model.parameters()) 164 | self.master_params = self.model_params 165 | self.param_groups_and_shapes = None 166 | self.lg_loss_scale = initial_lg_loss_scale 167 | 168 | # if self.use_fp16: 169 | self.param_groups_and_shapes = get_param_groups_and_shapes( 170 | self.model.named_parameters() 171 | ) 172 | self.master_params = make_master_params(self.param_groups_and_shapes) 173 | if self.use_fp16: 174 | self.model.convert_to_fp16() 175 | 176 | def zero_grad(self): 177 | zero_grad(self.model_params) 178 | 179 | def backward(self, loss: th.Tensor): 180 | if self.use_fp16: 181 | loss_scale = 2**self.lg_loss_scale 182 | (loss * loss_scale).backward() 183 | else: 184 | loss.backward() 185 | 186 | def optimize(self, opt: th.optim.Optimizer): 187 | if self.use_fp16: 188 | return self._optimize_fp16(opt) 189 | else: 190 | return self._optimize_normal(opt) 191 | 192 | def _optimize_fp16(self, opt: th.optim.Optimizer): 193 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 194 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 195 | grad_norm, param_norm = self._compute_norms(grad_scale=2**self.lg_loss_scale) 196 | if check_overflow(grad_norm): 197 | self.lg_loss_scale -= 1 198 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 199 | zero_master_grads(self.master_params) 200 | return False 201 | 202 | logger.logkv_mean("grad_norm", grad_norm) 203 | logger.logkv_mean("param_norm", param_norm) 204 | 205 | for p in self.master_params: 206 | p.grad.mul_(1.0 / (2**self.lg_loss_scale)) 207 | opt.step() 208 | zero_master_grads(self.master_params) 209 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 210 | self.lg_loss_scale += self.fp16_scale_growth 211 | return True 212 | 213 | def _optimize_normal(self, opt: th.optim.Optimizer): 214 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 215 | grad_norm, param_norm = self._compute_norms() 216 | if check_overflow(grad_norm) or np.isnan(grad_norm): 217 | logger.log(f"Found NaN; Skip current batch;") 218 | zero_master_grads(self.master_params) 219 | return False 220 | 221 | logger.logkv_mean("grad_norm", grad_norm) 222 | logger.logkv_mean("param_norm", param_norm) 223 | opt.step() 224 | zero_master_grads(self.master_params) 225 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 226 | return True 227 | 228 | def _compute_norms(self, grad_scale=1.0): 229 | grad_norm = 0.0 230 | param_norm = 0.0 231 | for p in self.master_params: 232 | with th.no_grad(): 233 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 234 | if p.grad is not None: 235 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 236 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 237 | 238 | def master_params_to_state_dict(self, master_params): 239 | return master_params_to_state_dict( 240 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 241 | ) 242 | 243 | def state_dict_to_master_params(self, state_dict): 244 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 245 | 246 | 247 | def check_overflow(value): 248 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 249 | -------------------------------------------------------------------------------- /BCM/cm/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 | if True: 100 | state_dict = model.state_dict() 101 | for master_param, (param_group, _) in zip( 102 | master_params, param_groups_and_shapes 103 | ): 104 | for (name, _), unflat_master_param in zip( 105 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 106 | ): 107 | assert name in state_dict 108 | state_dict[name] = unflat_master_param 109 | else: 110 | state_dict = model.state_dict() 111 | for i, (name, _value) in enumerate(model.named_parameters()): 112 | assert name in state_dict 113 | state_dict[name] = master_params[i] 114 | return state_dict 115 | 116 | 117 | def state_dict_to_master_params(model, state_dict, use_fp16): 118 | # if use_fp16: 119 | if True: 120 | named_model_params = [ 121 | (name, state_dict[name]) for name, _ in model.named_parameters() 122 | ] 123 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 124 | master_params = make_master_params(param_groups_and_shapes) 125 | else: 126 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 127 | return master_params 128 | 129 | 130 | def zero_master_grads(master_params): 131 | for param in master_params: 132 | param.grad = None 133 | 134 | 135 | def zero_grad(model_params): 136 | for param in model_params: 137 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 138 | if param.grad is not None: 139 | param.grad.detach_() 140 | param.grad.zero_() 141 | 142 | 143 | def param_grad_or_zeros(param): 144 | if param.grad is not None: 145 | return param.grad.data.detach() 146 | else: 147 | return th.zeros_like(param) 148 | 149 | 150 | class MixedPrecisionTrainer: 151 | def __init__( 152 | self, 153 | *, 154 | model, 155 | use_fp16=False, 156 | fp16_scale_growth=1e-3, 157 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 158 | ): 159 | self.model = model 160 | self.use_fp16 = use_fp16 161 | self.fp16_scale_growth = fp16_scale_growth 162 | 163 | self.model_params = list(self.model.parameters()) 164 | self.master_params = self.model_params 165 | self.param_groups_and_shapes = None 166 | self.lg_loss_scale = initial_lg_loss_scale 167 | 168 | # if self.use_fp16: 169 | self.param_groups_and_shapes = get_param_groups_and_shapes( 170 | self.model.named_parameters() 171 | ) 172 | self.master_params = make_master_params(self.param_groups_and_shapes) 173 | if self.use_fp16: 174 | self.model.convert_to_fp16() 175 | 176 | def zero_grad(self): 177 | zero_grad(self.model_params) 178 | 179 | def backward(self, loss: th.Tensor): 180 | if self.use_fp16: 181 | loss_scale = 2**self.lg_loss_scale 182 | (loss * loss_scale).backward() 183 | else: 184 | loss.backward() 185 | 186 | def optimize(self, opt: th.optim.Optimizer): 187 | if self.use_fp16: 188 | return self._optimize_fp16(opt) 189 | else: 190 | return self._optimize_normal(opt) 191 | 192 | def _optimize_fp16(self, opt: th.optim.Optimizer): 193 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 194 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 195 | grad_norm, param_norm = self._compute_norms(grad_scale=2**self.lg_loss_scale) 196 | if check_overflow(grad_norm): 197 | self.lg_loss_scale -= 1 198 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 199 | zero_master_grads(self.master_params) 200 | return False 201 | 202 | logger.logkv_mean("grad_norm", grad_norm) 203 | logger.logkv_mean("param_norm", param_norm) 204 | 205 | for p in self.master_params: 206 | p.grad.mul_(1.0 / (2**self.lg_loss_scale)) 207 | opt.step() 208 | zero_master_grads(self.master_params) 209 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 210 | self.lg_loss_scale += self.fp16_scale_growth 211 | return True 212 | 213 | def _optimize_normal(self, opt: th.optim.Optimizer): 214 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 215 | grad_norm, param_norm = self._compute_norms() 216 | if check_overflow(grad_norm) or np.isnan(grad_norm): 217 | logger.log(f"Found NaN; Skip current batch;") 218 | zero_master_grads(self.master_params) 219 | return False 220 | 221 | logger.logkv_mean("grad_norm", grad_norm) 222 | logger.logkv_mean("param_norm", param_norm) 223 | # 做一下grad clip 224 | nn.utils.clip_grad_norm_(self.master_params, 1.0) 225 | grad_norm, param_norm = self._compute_norms() 226 | logger.logkv_mean("clipped_grad_norm", grad_norm) 227 | 228 | opt.step() 229 | zero_master_grads(self.master_params) 230 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 231 | return True 232 | 233 | def _compute_norms(self, grad_scale=1.0): 234 | grad_norm = 0.0 235 | param_norm = 0.0 236 | for p in self.master_params: 237 | with th.no_grad(): 238 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 239 | if p.grad is not None: 240 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 241 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 242 | 243 | def master_params_to_state_dict(self, master_params): 244 | return master_params_to_state_dict( 245 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 246 | ) 247 | 248 | def state_dict_to_master_params(self, state_dict): 249 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 250 | 251 | 252 | def check_overflow(value): 253 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 254 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Bidirectional Consistency Models (PyTorch)
Official code and model checkpoints 2 | 3 | [![](http://img.shields.io/badge/arXiv-2403.18035-B31B1B.svg)](https://arxiv.org/abs/2403.18035) 4 | 5 | This repo contains: 6 | - official PyTorch **code** and model **weights** of [Bidirectional Consistency Models (BCM)](https://arxiv.org/abs/2403.18035) on ImageNet-64. 7 | - PyTorch **code** and model **weights** of our reproduced [Improved Consistency Training (iCT)](https://arxiv.org/abs/2310.14189) on ImageNet-64. 8 | 9 | BCM learns a single neural network that enables both forward and backward traversal along the PF ODE, efficiently unifying generation and inversion tasks within one framework. Our repository is based on [openai/consistency_models](https://github.com/openai/consistency_models), which was initially released under the MIT license. 10 | 11 | 12 | 13 | ## TL;DR 14 | BCM learns a single neural network that enables both forward and backward traversal along the PF ODE, efficiently unifying generation and inversion tasks within one framework. BCM offers diverse sampling options and has great potential in downstream tasks. 15 | 16 | ## Model Weights 17 | We provide checkpoints for BCM and our reproduced iCT on ImageNet-64: 18 | - [BCM-ImageNet-64](https://figshare.com/ndownloader/articles/27134694/versions/1?folder_path=bcf_imagenet64_no32_qkv_4096) 19 | - [iCT-ImageNet-64](https://figshare.com/ndownloader/articles/27134694/versions/1?folder_path=ict_imagenet64_no32_qkv_4096) 20 | - [BCM-deep-ImageNet-64](https://figshare.com/ndownloader/articles/27134694/versions/1?folder_path=bcf_imagenet64_no32_qkv_deep_4096) 21 | - [iCT-deep-ImageNet-64](https://figshare.com/ndownloader/articles/27134694/versions/2?folder_path=ict_imagenet64_no32_qkv_deep_4096) 22 | 23 | Their FIDs are as follows: 24 | | Name | NFE | FID | 25 | |:--------:|:---:|:---:| 26 | | BCM / BCM-deep | 1 | 4.18 / 3.14 | 27 | | | 2 | 2.88 / 2.45 | 28 | | | 3 | 2.78 / 2.61 | 29 | | | 4 | 2.68 / 2.35 | 30 | | reproduced iCT / iCT-deep| 1 | 4.60 / 3.94 | 31 | | | 2 | 3.40 / 3.14 | 32 | 33 | 34 | ## Dependencies 35 | 36 | To install all packages in this codebase along with their dependencies, run 37 | ``` 38 | cd iCT 39 | pip install -e . 40 | ``` 41 | 42 | To install with Docker, run the following commands: 43 | ``` 44 | cd docker && make build && make run 45 | ``` 46 | 47 | Please note that flash-attn==0.2.8, which cannot be substituted with the latest version and could be hard to install, is fortunately optional and not used in our best models. 48 | 49 | We also suggest manually install mpi4py using Anaconda instead of pip, with the following command: 50 | ``` 51 | conda install -c conda-forge mpi4py=3.1.4 mpich=3.2.3 52 | ``` 53 | 54 | ## Training 55 | 56 | As we described in our paper, for complex dataset like ImageNet-64, we propose to finetune BCM from pretrained iCT model. 57 | We, therefore, first provide code for iCT and then for BCM Finetuning. 58 | 59 | 60 | ### iCT 61 | 62 | The code for our reproduced iCT is located in the ```iCT``` folder. 63 | As we described in our paper, we found the original iCT suffers from instability on ImageNet-64. 64 | In our experiments, it diverges after ~620k iterations and the best one-step generation FID we got is ~6.20, largely falling behind the reported 4.02 in the iCT paper. 65 | We are open to any discussions on solutions to the instability issue and possible ways to reproduce the officially reported results. 66 | 67 | We suspect this instability comes from the architecture of ADM. Therefore, as a remedy, we proposed *removing the attention at the resolution of 32* and applying *normalization to QKV matrices*, following EDM2. We found this helpful in improving the performance and yielding a one-step FID of 4.60. 68 | We also apply *early stop* and save the checkpoint with the best one-step generation FID. 69 | 70 | Without modifications to the code, it is expected to start the training scripts with MPI for DDP training. For the commonly used SLURM, we provide the following starting script as an example: 71 | ``` 72 | srun -p YOUR_SLURM_PARTITION \ 73 | --job-name=ict_no32_qkv \ 74 | -n 64 --gres=gpu:8 --ntasks-per-node=8 \ 75 | --cpus-per-task=16 \ 76 | --quotatype=reserved \ 77 | --mpi=pmi2 \ 78 | sh WORKSPACE_DIR/iCT/scripts/ict_imagenet64_no32_qkv_4096.sh 79 | ``` 80 | The above script starts an iCT experiment with our architecture modifications, using 8 computing nodes (64 GPUs in total). 81 | 82 | To run the original iCT, please first switch back to the original network architecture. 83 | If you have flash-attn==0.2.8 installed, this can be done by simply setting ```attention_type="flash"``` at https://github.com/Mosasaur5526/BCM-iCT-torch/blob/main/iCT/cm/unet.py#L282. 84 | If not, just keep ```attention_type="default"``` and set ```cosine=False``` at https://github.com/Mosasaur5526/BCM-iCT-torch/blob/main/iCT/cm/unet.py#L412. 85 | Then run the following script: 86 | ``` 87 | srun -p YOUR_SLURM_PARTITION \ 88 | --job-name=ict \ 89 | -n 64 --gres=gpu:8 --ntasks-per-node=8 \ 90 | --cpus-per-task=16 \ 91 | --quotatype=reserved \ 92 | --mpi=pmi2 \ 93 | sh WORKSPACE_DIR/iCT/scripts/ict_imagenet64.sh 94 | ``` 95 | 96 | 97 | ### BCM Funetuning 98 | 99 | The code for BCM is located in the ```BCM``` folder. 100 | For ImageNet-64, we finetune BCM from pretrained iCT model to increase scalability, so please specify the location of the pretrained checkpoint in ```BCM/scripts/bcf_imagenet64_no32_qkv_4096.sh```. 101 | We carefully initialize the model to ensure that newly added ```t_end``` will not influence the iCT prediction. Please find the details in our paper. 102 | 103 | To perform BCF with, e.g., 64 GPUs, please run the following script: 104 | ``` 105 | srun -p YOUR_SLURM_PARTITION \ 106 | --job-name=bcm \ 107 | -n 64 --gres=gpu:8 --ntasks-per-node=8 \ 108 | --cpus-per-task=16 \ 109 | --quotatype=reserved \ 110 | --mpi=pmi2 \ 111 | sh WORKSPACE_DIR/BCM/scripts/bcf_imagenet64_no32_qkv_4096.sh 112 | ``` 113 | 114 | 115 | 116 | 117 | ### FP32 Training 118 | 119 | Our implementation also support training with fp32 by setting ```fp16=False``` in the training script, **which is actually *not* supported by [the official CM implementation](https://github.com/openai/consistency_models).** 120 | Please note that training with higher numerical accuracy doubles the computing budget and GPU memory and, according to our early experiments, may lead to different model behaviors during training. 121 | We hope our code and observation could help future studies on the influence of numerical issues on CMs. 122 | 123 | ## Evaluations 124 | 125 | ### Sampling 126 | 127 | Since BCM supports very flexible ways of sampling (ancestral, zigzag, mixture; see details in our paper), we think it would be overly verbose and less straightforward to pass arguments to the sampling script. 128 | Instead, we provide just one simple script (```BCM/scripts/image_sample.py``` or ```iCT/scripts/image_sample.py``` for BCM/iCT), and allow users to modify the code for all sampling methods. 129 | We provide detailed examples in the script, around https://github.com/Mosasaur5526/BCM-iCT-torch/blob/main/iCT/scripts/image_sample.py#L70 for iCT and around https://github.com/Mosasaur5526/BCM-iCT-torch/blob/main/BCM/scripts/image_sample.py#L116 for BCM. 130 | We believe these examples are simple and straightforward enough as each of them only requires to modify numbers in a few lines. 131 | 132 | To do distributed sampling on 4 GPUs (e.g., for iCT), please run: 133 | ``` 134 | srun -p YOUR_SLURM_PARTITION \ 135 | --job-name=ict_sampling \ 136 | -n 4 --gres=gpu:4 --ntasks-per-node=4 \ 137 | --cpus-per-task=16 \ 138 | --quotatype=reserved \ 139 | --mpi=pmi2 \ 140 | sh WORKSPACE_DIR/iCT/scripts/imagenet64_sample.sh 141 | ``` 142 | In the example script, it loads weights from ```CKPT_DIR/ict_imagenet64_no32_qkv_4096/ema_0.99997_680000.pt```, samples 50,000 images and saves them to ```WORKSPACE_DIR/samples/ict_imagenet64_no32_qkv_4096``` for further evaluation. 143 | 144 | ### Inversion and Reconstruction (BCM only) 145 | Inversion and reconstruction shares the same scripts as sampling. 146 | By setting ```--eval_mse=True``` in the sampling script, one can perform inversion and reconstruction for the images in ```--test_data_dir```. 147 | The per pixel MSE will be calculated automatically at the end and both the original and reconstructed images will be saved. 148 | Again for conciseness and clarity, we refer users to https://github.com/Mosasaur5526/BCM-iCT-torch/blob/main/BCM/scripts/image_sample.py#L172 to modify the code to enable one/multi-step inversion. 149 | 150 | Note that the ImageNet validation set is not structured by categories as the training set, so we modify the ```load_data``` function in ```cm/image_datasets.py ``` to support loading both images and labels from the valiadtion set. 151 | For convenience, the labels could be found in ```datasets/imagenet_val_label.txt``` and specified at https://github.com/Mosasaur5526/BCM-iCT-torch/blob/main/BCM/cm/image_datasets.py#L52; one may also load the image-label pairs in their customized ways by rewriting the loading function. 152 | Please notice the labels are important as they will be sent into the model as conditions during inversion and reconstruction. 153 | 154 | 155 | 156 | ### Calculating Metrics 157 | We follow the standard evaluation process in the [ADM repo](https://github.com/openai/guided-diffusion/tree/main/evaluations), as also adopted in the official CM repo. 158 | 159 | 160 | ### Visualizing Samples 161 | We also provide a simple visualization script in ```scripts/visualize_image.py```. 162 | 163 | ## Citation 164 | If you use this repository, including our code or the weights for BCM and our reproduced iCT, please cite the following work: 165 | ``` 166 | @article{li2024bidirectional, 167 | title={Bidirectional Consistency Models}, 168 | author={Li, Liangchen and He, Jiajun}, 169 | journal={arXiv preprint arXiv:2403.18035}, 170 | year={2024} 171 | } 172 | ``` 173 | -------------------------------------------------------------------------------- /BCM/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 | import sys 6 | sys.path.append('../BCM') 7 | 8 | from PIL import Image 9 | 10 | import argparse 11 | import os 12 | 13 | import numpy as np 14 | import torch as th 15 | import torch.distributed as dist 16 | 17 | from cm import dist_util, logger 18 | from cm.script_util import ( 19 | NUM_CLASSES, 20 | model_and_diffusion_defaults, 21 | create_model_and_diffusion, 22 | add_dict_to_argparser, 23 | args_to_dict, 24 | ) 25 | from cm.random_util import get_generator 26 | from cm.karras_diffusion import karras_sample 27 | from cm.image_datasets import load_data 28 | 29 | 30 | def center_crop_arr(pil_image, image_size): 31 | # We are not on a new enough PIL to support the `reducing_gap` 32 | # argument, which uses BOX downsampling at powers of two first. 33 | # Thus, we do it by hand to improve downsample quality. 34 | while min(*pil_image.size) >= 2 * image_size: 35 | pil_image = pil_image.resize( 36 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 37 | ) 38 | 39 | scale = image_size / min(*pil_image.size) 40 | pil_image = pil_image.resize( 41 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 42 | ) 43 | 44 | arr = np.array(pil_image) 45 | crop_y = (arr.shape[0] - image_size) // 2 46 | crop_x = (arr.shape[1] - image_size) // 2 47 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 48 | 49 | 50 | def main(): 51 | args = create_argparser().parse_args() 52 | 53 | dist_util.setup_dist() 54 | logger.configure(dir=os.path.join(args.save_dir, args.exp_name)) 55 | 56 | if "consistency" in args.training_mode: 57 | distillation = True 58 | else: 59 | distillation = False 60 | 61 | if args.eval_mse: 62 | logger.log("loading test data for reconstruction MSE evaluation...") 63 | # test data 64 | data = load_data( 65 | data_dir=args.test_data_dir, 66 | batch_size=args.batch_size, 67 | image_size=args.image_size, 68 | class_cond=args.class_cond, 69 | val=True 70 | ) 71 | else: 72 | data = None 73 | 74 | logger.log("creating model and diffusion...") 75 | model, diffusion = create_model_and_diffusion( 76 | **args_to_dict(args, model_and_diffusion_defaults().keys()), 77 | distillation=distillation, 78 | model_type=args.model_type 79 | ) 80 | model.load_state_dict( 81 | dist_util.load_state_dict(args.model_path, map_location="cpu") 82 | ) 83 | model.to(dist_util.dev()) 84 | if args.use_fp16: 85 | model.convert_to_fp16() 86 | model.eval() 87 | model_path = args.model_path 88 | model_path = model_path.split('/')[-1][:-3] 89 | 90 | logger.log("sampling...") 91 | if args.sampler == "multistep": 92 | assert len(args.ts) > 0 93 | ts = tuple(int(x) for x in args.ts.split(",")) 94 | else: 95 | ts = None 96 | 97 | all_images = [] 98 | all_images2 = [] 99 | all_labels = [] 100 | generator = get_generator(args.generator, args.num_samples, args.seed) 101 | 102 | while len(all_images) * args.batch_size < args.num_samples: 103 | if not args.eval_mse: 104 | # then perform generation 105 | model_kwargs = {} 106 | if args.class_cond: 107 | classes = th.randint( 108 | low=0, high=NUM_CLASSES, size=(args.batch_size,), device=dist_util.dev() 109 | ) 110 | model_kwargs["y"] = classes 111 | 112 | # now, sample is the initial Gaussian noise at T=80.0 113 | sample = th.randn((args.batch_size, 3, args.image_size, args.image_size), device=dist_util.dev()) * 80.0 114 | multiplier = th.ones(sample.shape[0], dtype=sample.dtype, device=sample.device) 115 | 116 | # ----------------- one-step generation ----------------- 117 | # 80.0 -> 0.002 118 | _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 80.0, sigmas_end=multiplier * 0.002, y=classes) 119 | # ---------------------------------------------------------- 120 | 121 | # -------------- two-step ancestral sampling ------------ 122 | # 80.0 -> 2.4 -> 0.002 123 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 80.0, sigmas_end=multiplier * 2.4, y=classes) 124 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 2.4, sigmas_end=multiplier * 0.002, y=classes) 125 | # ---------------------------------------------------------- 126 | 127 | # -------------- three-step zigzag sampling ------------ 128 | # 80.0 -> 0.002 -> 0.1 (manually added noise) -> 1.2 (amplified by BCM) -> 0.002 129 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 80.0, sigmas_end=multiplier * 0.002, y=classes) 130 | # sample += th.randn_like(sample, device=sample.device) * 0.1 131 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 0.1, sigmas_end=multiplier * 1.2, y=classes) 132 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 1.2, sigmas_end=multiplier * 0.002, y=classes) 133 | # ---------------------------------------------------------- 134 | 135 | # -------------- four-step mixed sampling -------------- 136 | # 80.0 -> 3.0 -> 0.002 137 | # -> 0.12 (manually added noise) -> 0.4 (amplified by BCM) -> 0.002 138 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 80.0, sigmas_end=multiplier * 3.0, y=classes) 139 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 3.0, sigmas_end=multiplier * 0.002, y=classes) 140 | # sample += th.randn_like(sample, device=sample.device) * 0.12 141 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 0.12, sigmas_end=multiplier * 0.4, y=classes) 142 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 0.4, sigmas_end=multiplier * 0.002, y=classes) 143 | # ---------------------------------------------------------- 144 | 145 | else: 146 | # to perform inversion and evaluate reconstrution MSE, first load samples from val set 147 | batch, classes = next(data) 148 | batch = batch.to(device=dist_util.dev()) 149 | classes = classes['y'].to(device=dist_util.dev()) 150 | sample = batch 151 | 152 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 153 | sample = sample.permute(0, 2, 3, 1) 154 | sample = sample.contiguous() 155 | 156 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 157 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 158 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 159 | if args.class_cond: 160 | gathered_labels = [ 161 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 162 | ] 163 | dist.all_gather(gathered_labels, classes) 164 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 165 | logger.log(f"created {len(all_images) * args.batch_size} samples") 166 | 167 | if args.reconstruct or args.eval_mse: 168 | # add a small initial noise 169 | multiplier = th.ones(sample.shape[0], dtype=sample.dtype, device=sample.device) 170 | sample = sample + th.randn_like(sample, device=sample.device) * 0.07 171 | 172 | # ------------------- one-step inversion ----------------- 173 | # _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 0.07, sigmas_end=multiplier * 80.0, y=classes) 174 | # ---------------------------------------------------------- 175 | 176 | # ------------------- two-step inversion ----------------- 177 | _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 0.07, sigmas_end=multiplier * 15.0, y=classes) 178 | _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 15.0, sigmas_end=multiplier * 80.0, y=classes) 179 | # ---------------------------------------------------------- 180 | 181 | # --------------------- reconstrution -------------------- 182 | _, sample = diffusion.denoise(model, sample, sigmas=multiplier * 80.0, sigmas_end=multiplier * 0.002, y=classes) 183 | # ---------------------------------------------------------- 184 | 185 | sample = sample.permute(0, 2, 3, 1) 186 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 187 | sample = sample.contiguous() 188 | 189 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 190 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 191 | all_images2.extend([sample.cpu().numpy() for sample in gathered_samples]) 192 | 193 | arr = np.concatenate(all_images, axis=0) 194 | arr = arr[: args.num_samples] 195 | if args.reconstruct or args.eval_mse: 196 | arr2 = np.concatenate(all_images2, axis=0) 197 | arr2 = arr2[: args.num_samples] 198 | if args.class_cond: 199 | label_arr = np.concatenate(all_labels, axis=0) 200 | label_arr = label_arr[: args.num_samples] 201 | if dist.get_rank() == 0: 202 | shape_str = "x".join([str(x) for x in arr.shape]) 203 | name = 'samples_original' if args.reconstruct or args.eval_mse else 'samples' 204 | name2 = 'Original samples' if args.reconstruct or args.eval_mse else 'Samples' 205 | out_path = os.path.join(logger.get_dir(), f"{model_path}_{name}_{shape_str}.npz") 206 | logger.log(f"{name2} saving to {out_path}") 207 | if args.class_cond: 208 | np.savez(out_path, arr, label_arr) 209 | else: 210 | np.savez(out_path, arr) 211 | if args.reconstruct or args.eval_mse: 212 | out_path2 = os.path.join(logger.get_dir(), f"{model_path}_samples_reconstructed.npz") 213 | np.savez(out_path2, arr2) 214 | logger.log(f"Reconstructed images saving to {out_path2}") 215 | 216 | mse = (((arr / 255. - arr2 / 255.).reshape(-1)) ** 2).mean() 217 | logger.log(f"Reconstruction per pixel MSE: {mse}") 218 | 219 | dist.barrier() 220 | logger.log("sampling complete") 221 | 222 | 223 | def create_argparser(): 224 | defaults = dict( 225 | training_mode="edm", 226 | generator="determ", 227 | clip_denoised=True, 228 | num_samples=50000, 229 | batch_size=16, 230 | sampler="heun", 231 | s_churn=0.0, 232 | s_tmin=0.0, 233 | s_tmax=float("inf"), 234 | s_noise=1.0, 235 | steps=40, 236 | model_path="", 237 | save_dir='./checkpoints', 238 | exp_name='ict', 239 | seed=42, 240 | ts="", 241 | reconstruct=False, 242 | test_data_dir='/mnt/petrelfs/share/images/val', 243 | eval_mse=False, 244 | model_type='tsinghua' 245 | ) 246 | defaults.update(model_and_diffusion_defaults()) 247 | parser = argparse.ArgumentParser() 248 | add_dict_to_argparser(parser, defaults) 249 | return parser 250 | 251 | 252 | if __name__ == "__main__": 253 | main() 254 | -------------------------------------------------------------------------------- /BCM/evaluations/th_evaluator.py: -------------------------------------------------------------------------------- 1 | from .inception_v3 import InceptionV3 2 | import blobfile as bf 3 | import torch 4 | import torch.distributed as dist 5 | import torch.nn as nn 6 | from cm import dist_util 7 | import numpy as np 8 | import warnings 9 | from scipy import linalg 10 | from PIL import Image 11 | from tqdm import tqdm 12 | 13 | 14 | def clip_preproc(preproc_fn, x): 15 | return preproc_fn(Image.fromarray(x.astype(np.uint8))) 16 | 17 | 18 | def all_gather(x, dim=0): 19 | xs = [torch.zeros_like(x) for _ in range(dist.get_world_size())] 20 | dist.all_gather(xs, x) 21 | return torch.cat(xs, dim=dim) 22 | 23 | 24 | class FIDStatistics: 25 | def __init__(self, mu: np.ndarray, sigma: np.ndarray, resolution: int): 26 | self.mu = mu 27 | self.sigma = sigma 28 | self.resolution = resolution 29 | 30 | def frechet_distance(self, other, eps=1e-6): 31 | """ 32 | Compute the Frechet distance between two sets of statistics. 33 | """ 34 | mu1, sigma1 = self.mu, self.sigma 35 | mu2, sigma2 = other.mu, other.sigma 36 | 37 | mu1 = np.atleast_1d(mu1) 38 | mu2 = np.atleast_1d(mu2) 39 | 40 | sigma1 = np.atleast_2d(sigma1) 41 | sigma2 = np.atleast_2d(sigma2) 42 | 43 | assert ( 44 | mu1.shape == mu2.shape 45 | ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}" 46 | assert ( 47 | sigma1.shape == sigma2.shape 48 | ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}" 49 | 50 | diff = mu1 - mu2 51 | 52 | # product might be almost singular 53 | covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) 54 | if not np.isfinite(covmean).all(): 55 | msg = ( 56 | "fid calculation produces singular product; adding %s to diagonal of cov estimates" 57 | % eps 58 | ) 59 | warnings.warn(msg) 60 | offset = np.eye(sigma1.shape[0]) * eps 61 | covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) 62 | 63 | # numerical error might give slight imaginary component 64 | if np.iscomplexobj(covmean): 65 | if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): 66 | m = np.max(np.abs(covmean.imag)) 67 | raise ValueError("Imaginary component {}".format(m)) 68 | covmean = covmean.real 69 | 70 | tr_covmean = np.trace(covmean) 71 | 72 | return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean 73 | 74 | 75 | class FIDAndIS: 76 | def __init__( 77 | self, 78 | softmax_batch_size=512, 79 | clip_score_batch_size=512, 80 | path="https://openaipublic.blob.core.windows.net/consistency/inception/inception-2015-12-05.pt", 81 | ): 82 | import clip 83 | 84 | super().__init__() 85 | 86 | self.softmax_batch_size = softmax_batch_size 87 | self.clip_score_batch_size = clip_score_batch_size 88 | self.inception = InceptionV3() 89 | with bf.BlobFile(path, "rb") as f: 90 | self.inception.load_state_dict(torch.load(f)) 91 | self.inception.eval() 92 | self.inception.to(dist_util.dev()) 93 | 94 | self.inception_softmax = self.inception.create_softmax_model() 95 | 96 | if dist.get_rank() % 8 == 0: 97 | clip_model, self.clip_preproc_fn = clip.load( 98 | "ViT-B/32", device=dist_util.dev() 99 | ) 100 | dist.barrier() 101 | if dist.get_rank() % 8 != 0: 102 | clip_model, self.clip_preproc_fn = clip.load( 103 | "ViT-B/32", device=dist_util.dev() 104 | ) 105 | dist.barrier() 106 | 107 | # Compute the probe features separately from the final projection. 108 | class ProjLayer(nn.Module): 109 | def __init__(self, param): 110 | super().__init__() 111 | self.param = param 112 | 113 | def forward(self, x): 114 | return x @ self.param 115 | 116 | self.clip_visual = clip_model.visual 117 | self.clip_proj = ProjLayer(self.clip_visual.proj) 118 | self.clip_visual.proj = None 119 | 120 | class TextModel(nn.Module): 121 | def __init__(self, clip_model): 122 | super().__init__() 123 | self.clip_model = clip_model 124 | 125 | def forward(self, x): 126 | return self.clip_model.encode_text(x) 127 | 128 | self.clip_tokenizer = lambda captions: clip.tokenize(captions, truncate=True) 129 | self.clip_text = TextModel(clip_model) 130 | self.clip_logit_scale = clip_model.logit_scale.exp().item() 131 | self.ref_features = {} 132 | self.is_root = not dist.is_initialized() or dist.get_rank() == 0 133 | 134 | def get_statistics(self, activations: np.ndarray, resolution: int): 135 | """ 136 | Compute activation statistics for a batch of images. 137 | 138 | :param activations: an [N x D] batch of activations. 139 | :return: an FIDStatistics object. 140 | """ 141 | mu = np.mean(activations, axis=0) 142 | sigma = np.cov(activations, rowvar=False) 143 | return FIDStatistics(mu, sigma, resolution) 144 | 145 | def get_preds(self, batch, captions=None): 146 | with torch.no_grad(): 147 | batch = 127.5 * (batch + 1) 148 | np_batch = batch.to(torch.uint8).cpu().numpy().transpose((0, 2, 3, 1)) 149 | 150 | pred, spatial_pred = self.inception(batch) 151 | pred, spatial_pred = pred.reshape( 152 | [pred.shape[0], -1] 153 | ), spatial_pred.reshape([spatial_pred.shape[0], -1]) 154 | 155 | clip_in = torch.stack( 156 | [clip_preproc(self.clip_preproc_fn, img) for img in np_batch] 157 | ) 158 | clip_pred = self.clip_visual(clip_in.half().to(dist_util.dev())) 159 | if captions is not None: 160 | text_in = self.clip_tokenizer(captions) 161 | text_pred = self.clip_text(text_in.to(dist_util.dev())) 162 | else: 163 | # Hack to easily deal with no captions 164 | text_pred = self.clip_proj(clip_pred.half()) 165 | text_pred = text_pred / text_pred.norm(dim=-1, keepdim=True) 166 | 167 | return pred, spatial_pred, clip_pred, text_pred, np_batch 168 | 169 | def get_inception_score( 170 | self, activations: np.ndarray, split_size: int = 5000 171 | ) -> float: 172 | """ 173 | Compute the inception score using a batch of activations. 174 | :param activations: an [N x D] batch of activations. 175 | :param split_size: the number of samples per split. This is used to 176 | make results consistent with other work, even when 177 | using a different number of samples. 178 | :return: an inception score estimate. 179 | """ 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 | with torch.no_grad(): 184 | softmax_out.append( 185 | self.inception_softmax(torch.from_numpy(acts).to(dist_util.dev())) 186 | .cpu() 187 | .numpy() 188 | ) 189 | preds = np.concatenate(softmax_out, axis=0) 190 | # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46 191 | scores = [] 192 | for i in range(0, len(preds), split_size): 193 | part = preds[i : i + split_size] 194 | kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))) 195 | kl = np.mean(np.sum(kl, 1)) 196 | scores.append(np.exp(kl)) 197 | return float(np.mean(scores)) 198 | 199 | def get_clip_score( 200 | self, activations: np.ndarray, text_features: np.ndarray 201 | ) -> float: 202 | # Sizes should never mismatch, but if they do we want to compute 203 | # _some_ value instead of crash looping. 204 | size = min(len(activations), len(text_features)) 205 | activations = activations[:size] 206 | text_features = text_features[:size] 207 | 208 | scores_out = [] 209 | for i in range(0, len(activations), self.clip_score_batch_size): 210 | acts = activations[i : i + self.clip_score_batch_size] 211 | sub_features = text_features[i : i + self.clip_score_batch_size] 212 | with torch.no_grad(): 213 | image_features = self.clip_proj( 214 | torch.from_numpy(acts).half().to(dist_util.dev()) 215 | ) 216 | image_features = image_features / image_features.norm( 217 | dim=-1, keepdim=True 218 | ) 219 | image_features = image_features.detach().cpu().float().numpy() 220 | scores_out.extend(np.sum(sub_features * image_features, axis=-1).tolist()) 221 | return np.mean(scores_out) * self.clip_logit_scale 222 | 223 | def get_activations(self, data, num_samples, global_batch_size, pr_samples=50000): 224 | if self.is_root: 225 | preds = [] 226 | spatial_preds = [] 227 | clip_preds = [] 228 | pr_images = [] 229 | 230 | for _ in tqdm(range(0, int(np.ceil(num_samples / global_batch_size)))): 231 | batch, cond, _ = next(data) 232 | batch, cond = batch.to(dist_util.dev()), { 233 | k: v.to(dist_util.dev()) for k, v in cond.items() 234 | } 235 | pred, spatial_pred, clip_pred, _, np_batch = self.get_preds(batch) 236 | pred, spatial_pred, clip_pred = ( 237 | all_gather(pred).cpu().numpy(), 238 | all_gather(spatial_pred).cpu().numpy(), 239 | all_gather(clip_pred).cpu().numpy(), 240 | ) 241 | if self.is_root: 242 | preds.append(pred) 243 | spatial_preds.append(spatial_pred) 244 | clip_preds.append(clip_pred) 245 | if len(pr_images) * np_batch.shape[0] < pr_samples: 246 | pr_images.append(np_batch) 247 | 248 | if self.is_root: 249 | preds, spatial_preds, clip_preds, pr_images = ( 250 | np.concatenate(preds, axis=0), 251 | np.concatenate(spatial_preds, axis=0), 252 | np.concatenate(clip_preds, axis=0), 253 | np.concatenate(pr_images, axis=0), 254 | ) 255 | # assert len(pr_images) >= pr_samples 256 | return ( 257 | preds[:num_samples], 258 | spatial_preds[:num_samples], 259 | clip_preds[:num_samples], 260 | pr_images[:pr_samples], 261 | ) 262 | else: 263 | return [], [], [], [] 264 | 265 | def get_virtual_batch(self, data, num_samples, global_batch_size, resolution): 266 | preds, spatial_preds, clip_preds, batch = self.get_activations( 267 | data, num_samples, global_batch_size, pr_samples=10000 268 | ) 269 | if self.is_root: 270 | fid_stats = self.get_statistics(preds, resolution) 271 | spatial_stats = self.get_statistics(spatial_preds, resolution) 272 | clip_stats = self.get_statistics(clip_preds, resolution) 273 | return batch, dict( 274 | mu=fid_stats.mu, 275 | sigma=fid_stats.sigma, 276 | mu_s=spatial_stats.mu, 277 | sigma_s=spatial_stats.sigma, 278 | mu_clip=clip_stats.mu, 279 | sigma_clip=clip_stats.sigma, 280 | ) 281 | else: 282 | return None, dict() 283 | 284 | def set_ref_batch(self, ref_batch): 285 | with bf.BlobFile(ref_batch, "rb") as f: 286 | data = np.load(f) 287 | fid_stats = FIDStatistics(mu=data["mu"], sigma=data["sigma"], resolution=-1) 288 | spatial_stats = FIDStatistics( 289 | mu=data["mu_s"], sigma=data["sigma_s"], resolution=-1 290 | ) 291 | clip_stats = FIDStatistics( 292 | mu=data["mu_clip"], sigma=data["sigma_clip"], resolution=-1 293 | ) 294 | 295 | self.ref_features[ref_batch] = (fid_stats, spatial_stats, clip_stats) 296 | 297 | def get_ref_batch(self, ref_batch): 298 | return self.ref_features[ref_batch] 299 | -------------------------------------------------------------------------------- /iCT/evaluations/th_evaluator.py: -------------------------------------------------------------------------------- 1 | from .inception_v3 import InceptionV3 2 | import blobfile as bf 3 | import torch 4 | import torch.distributed as dist 5 | import torch.nn as nn 6 | from cm import dist_util 7 | import numpy as np 8 | import warnings 9 | from scipy import linalg 10 | from PIL import Image 11 | from tqdm import tqdm 12 | 13 | 14 | def clip_preproc(preproc_fn, x): 15 | return preproc_fn(Image.fromarray(x.astype(np.uint8))) 16 | 17 | 18 | def all_gather(x, dim=0): 19 | xs = [torch.zeros_like(x) for _ in range(dist.get_world_size())] 20 | dist.all_gather(xs, x) 21 | return torch.cat(xs, dim=dim) 22 | 23 | 24 | class FIDStatistics: 25 | def __init__(self, mu: np.ndarray, sigma: np.ndarray, resolution: int): 26 | self.mu = mu 27 | self.sigma = sigma 28 | self.resolution = resolution 29 | 30 | def frechet_distance(self, other, eps=1e-6): 31 | """ 32 | Compute the Frechet distance between two sets of statistics. 33 | """ 34 | mu1, sigma1 = self.mu, self.sigma 35 | mu2, sigma2 = other.mu, other.sigma 36 | 37 | mu1 = np.atleast_1d(mu1) 38 | mu2 = np.atleast_1d(mu2) 39 | 40 | sigma1 = np.atleast_2d(sigma1) 41 | sigma2 = np.atleast_2d(sigma2) 42 | 43 | assert ( 44 | mu1.shape == mu2.shape 45 | ), f"Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}" 46 | assert ( 47 | sigma1.shape == sigma2.shape 48 | ), f"Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}" 49 | 50 | diff = mu1 - mu2 51 | 52 | # product might be almost singular 53 | covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False) 54 | if not np.isfinite(covmean).all(): 55 | msg = ( 56 | "fid calculation produces singular product; adding %s to diagonal of cov estimates" 57 | % eps 58 | ) 59 | warnings.warn(msg) 60 | offset = np.eye(sigma1.shape[0]) * eps 61 | covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset)) 62 | 63 | # numerical error might give slight imaginary component 64 | if np.iscomplexobj(covmean): 65 | if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3): 66 | m = np.max(np.abs(covmean.imag)) 67 | raise ValueError("Imaginary component {}".format(m)) 68 | covmean = covmean.real 69 | 70 | tr_covmean = np.trace(covmean) 71 | 72 | return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean 73 | 74 | 75 | class FIDAndIS: 76 | def __init__( 77 | self, 78 | softmax_batch_size=512, 79 | clip_score_batch_size=512, 80 | path="https://openaipublic.blob.core.windows.net/consistency/inception/inception-2015-12-05.pt", 81 | ): 82 | import clip 83 | 84 | super().__init__() 85 | 86 | self.softmax_batch_size = softmax_batch_size 87 | self.clip_score_batch_size = clip_score_batch_size 88 | self.inception = InceptionV3() 89 | with bf.BlobFile(path, "rb") as f: 90 | self.inception.load_state_dict(torch.load(f)) 91 | self.inception.eval() 92 | self.inception.to(dist_util.dev()) 93 | 94 | self.inception_softmax = self.inception.create_softmax_model() 95 | 96 | if dist.get_rank() % 8 == 0: 97 | clip_model, self.clip_preproc_fn = clip.load( 98 | "ViT-B/32", device=dist_util.dev() 99 | ) 100 | dist.barrier() 101 | if dist.get_rank() % 8 != 0: 102 | clip_model, self.clip_preproc_fn = clip.load( 103 | "ViT-B/32", device=dist_util.dev() 104 | ) 105 | dist.barrier() 106 | 107 | # Compute the probe features separately from the final projection. 108 | class ProjLayer(nn.Module): 109 | def __init__(self, param): 110 | super().__init__() 111 | self.param = param 112 | 113 | def forward(self, x): 114 | return x @ self.param 115 | 116 | self.clip_visual = clip_model.visual 117 | self.clip_proj = ProjLayer(self.clip_visual.proj) 118 | self.clip_visual.proj = None 119 | 120 | class TextModel(nn.Module): 121 | def __init__(self, clip_model): 122 | super().__init__() 123 | self.clip_model = clip_model 124 | 125 | def forward(self, x): 126 | return self.clip_model.encode_text(x) 127 | 128 | self.clip_tokenizer = lambda captions: clip.tokenize(captions, truncate=True) 129 | self.clip_text = TextModel(clip_model) 130 | self.clip_logit_scale = clip_model.logit_scale.exp().item() 131 | self.ref_features = {} 132 | self.is_root = not dist.is_initialized() or dist.get_rank() == 0 133 | 134 | def get_statistics(self, activations: np.ndarray, resolution: int): 135 | """ 136 | Compute activation statistics for a batch of images. 137 | 138 | :param activations: an [N x D] batch of activations. 139 | :return: an FIDStatistics object. 140 | """ 141 | mu = np.mean(activations, axis=0) 142 | sigma = np.cov(activations, rowvar=False) 143 | return FIDStatistics(mu, sigma, resolution) 144 | 145 | def get_preds(self, batch, captions=None): 146 | with torch.no_grad(): 147 | batch = 127.5 * (batch + 1) 148 | np_batch = batch.to(torch.uint8).cpu().numpy().transpose((0, 2, 3, 1)) 149 | 150 | pred, spatial_pred = self.inception(batch) 151 | pred, spatial_pred = pred.reshape( 152 | [pred.shape[0], -1] 153 | ), spatial_pred.reshape([spatial_pred.shape[0], -1]) 154 | 155 | clip_in = torch.stack( 156 | [clip_preproc(self.clip_preproc_fn, img) for img in np_batch] 157 | ) 158 | clip_pred = self.clip_visual(clip_in.half().to(dist_util.dev())) 159 | if captions is not None: 160 | text_in = self.clip_tokenizer(captions) 161 | text_pred = self.clip_text(text_in.to(dist_util.dev())) 162 | else: 163 | # Hack to easily deal with no captions 164 | text_pred = self.clip_proj(clip_pred.half()) 165 | text_pred = text_pred / text_pred.norm(dim=-1, keepdim=True) 166 | 167 | return pred, spatial_pred, clip_pred, text_pred, np_batch 168 | 169 | def get_inception_score( 170 | self, activations: np.ndarray, split_size: int = 5000 171 | ) -> float: 172 | """ 173 | Compute the inception score using a batch of activations. 174 | :param activations: an [N x D] batch of activations. 175 | :param split_size: the number of samples per split. This is used to 176 | make results consistent with other work, even when 177 | using a different number of samples. 178 | :return: an inception score estimate. 179 | """ 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 | with torch.no_grad(): 184 | softmax_out.append( 185 | self.inception_softmax(torch.from_numpy(acts).to(dist_util.dev())) 186 | .cpu() 187 | .numpy() 188 | ) 189 | preds = np.concatenate(softmax_out, axis=0) 190 | # https://github.com/openai/improved-gan/blob/4f5d1ec5c16a7eceb206f42bfc652693601e1d5c/inception_score/model.py#L46 191 | scores = [] 192 | for i in range(0, len(preds), split_size): 193 | part = preds[i : i + split_size] 194 | kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))) 195 | kl = np.mean(np.sum(kl, 1)) 196 | scores.append(np.exp(kl)) 197 | return float(np.mean(scores)) 198 | 199 | def get_clip_score( 200 | self, activations: np.ndarray, text_features: np.ndarray 201 | ) -> float: 202 | # Sizes should never mismatch, but if they do we want to compute 203 | # _some_ value instead of crash looping. 204 | size = min(len(activations), len(text_features)) 205 | activations = activations[:size] 206 | text_features = text_features[:size] 207 | 208 | scores_out = [] 209 | for i in range(0, len(activations), self.clip_score_batch_size): 210 | acts = activations[i : i + self.clip_score_batch_size] 211 | sub_features = text_features[i : i + self.clip_score_batch_size] 212 | with torch.no_grad(): 213 | image_features = self.clip_proj( 214 | torch.from_numpy(acts).half().to(dist_util.dev()) 215 | ) 216 | image_features = image_features / image_features.norm( 217 | dim=-1, keepdim=True 218 | ) 219 | image_features = image_features.detach().cpu().float().numpy() 220 | scores_out.extend(np.sum(sub_features * image_features, axis=-1).tolist()) 221 | return np.mean(scores_out) * self.clip_logit_scale 222 | 223 | def get_activations(self, data, num_samples, global_batch_size, pr_samples=50000): 224 | if self.is_root: 225 | preds = [] 226 | spatial_preds = [] 227 | clip_preds = [] 228 | pr_images = [] 229 | 230 | for _ in tqdm(range(0, int(np.ceil(num_samples / global_batch_size)))): 231 | batch, cond, _ = next(data) 232 | batch, cond = batch.to(dist_util.dev()), { 233 | k: v.to(dist_util.dev()) for k, v in cond.items() 234 | } 235 | pred, spatial_pred, clip_pred, _, np_batch = self.get_preds(batch) 236 | pred, spatial_pred, clip_pred = ( 237 | all_gather(pred).cpu().numpy(), 238 | all_gather(spatial_pred).cpu().numpy(), 239 | all_gather(clip_pred).cpu().numpy(), 240 | ) 241 | if self.is_root: 242 | preds.append(pred) 243 | spatial_preds.append(spatial_pred) 244 | clip_preds.append(clip_pred) 245 | if len(pr_images) * np_batch.shape[0] < pr_samples: 246 | pr_images.append(np_batch) 247 | 248 | if self.is_root: 249 | preds, spatial_preds, clip_preds, pr_images = ( 250 | np.concatenate(preds, axis=0), 251 | np.concatenate(spatial_preds, axis=0), 252 | np.concatenate(clip_preds, axis=0), 253 | np.concatenate(pr_images, axis=0), 254 | ) 255 | # assert len(pr_images) >= pr_samples 256 | return ( 257 | preds[:num_samples], 258 | spatial_preds[:num_samples], 259 | clip_preds[:num_samples], 260 | pr_images[:pr_samples], 261 | ) 262 | else: 263 | return [], [], [], [] 264 | 265 | def get_virtual_batch(self, data, num_samples, global_batch_size, resolution): 266 | preds, spatial_preds, clip_preds, batch = self.get_activations( 267 | data, num_samples, global_batch_size, pr_samples=10000 268 | ) 269 | if self.is_root: 270 | fid_stats = self.get_statistics(preds, resolution) 271 | spatial_stats = self.get_statistics(spatial_preds, resolution) 272 | clip_stats = self.get_statistics(clip_preds, resolution) 273 | return batch, dict( 274 | mu=fid_stats.mu, 275 | sigma=fid_stats.sigma, 276 | mu_s=spatial_stats.mu, 277 | sigma_s=spatial_stats.sigma, 278 | mu_clip=clip_stats.mu, 279 | sigma_clip=clip_stats.sigma, 280 | ) 281 | else: 282 | return None, dict() 283 | 284 | def set_ref_batch(self, ref_batch): 285 | with bf.BlobFile(ref_batch, "rb") as f: 286 | data = np.load(f) 287 | fid_stats = FIDStatistics(mu=data["mu"], sigma=data["sigma"], resolution=-1) 288 | spatial_stats = FIDStatistics( 289 | mu=data["mu_s"], sigma=data["sigma_s"], resolution=-1 290 | ) 291 | clip_stats = FIDStatistics( 292 | mu=data["mu_clip"], sigma=data["sigma_clip"], resolution=-1 293 | ) 294 | 295 | self.ref_features[ref_batch] = (fid_stats, spatial_stats, clip_stats) 296 | 297 | def get_ref_batch(self, ref_batch): 298 | return self.ref_features[ref_batch] 299 | -------------------------------------------------------------------------------- /BCM/evaluations/inception_v3.py: -------------------------------------------------------------------------------- 1 | # Ported from the model here: 2 | # https://github.com/NVlabs/stylegan3/blob/407db86e6fe432540a22515310188288687858fa/metrics/frechet_inception_distance.py#L22 3 | # 4 | # I have verified that the spatial features and output features are correct 5 | # within a mean absolute error of ~3e-5. 6 | 7 | import collections 8 | 9 | import torch 10 | 11 | 12 | class Conv2dLayer(torch.nn.Module): 13 | def __init__(self, in_channels, out_channels, kh, kw, stride=1, padding=0): 14 | super().__init__() 15 | self.stride = stride 16 | self.padding = padding 17 | self.weight = torch.nn.Parameter(torch.zeros(out_channels, in_channels, kh, kw)) 18 | self.beta = torch.nn.Parameter(torch.zeros(out_channels)) 19 | self.mean = torch.nn.Parameter(torch.zeros(out_channels)) 20 | self.var = torch.nn.Parameter(torch.zeros(out_channels)) 21 | 22 | def forward(self, x): 23 | x = torch.nn.functional.conv2d( 24 | x, self.weight.to(x.dtype), stride=self.stride, padding=self.padding 25 | ) 26 | x = torch.nn.functional.batch_norm( 27 | x, running_mean=self.mean, running_var=self.var, bias=self.beta, eps=1e-3 28 | ) 29 | x = torch.nn.functional.relu(x) 30 | return x 31 | 32 | 33 | # ---------------------------------------------------------------------------- 34 | 35 | 36 | class InceptionA(torch.nn.Module): 37 | def __init__(self, in_channels, tmp_channels): 38 | super().__init__() 39 | self.conv = Conv2dLayer(in_channels, 64, kh=1, kw=1) 40 | self.tower = torch.nn.Sequential( 41 | collections.OrderedDict( 42 | [ 43 | ("conv", Conv2dLayer(in_channels, 48, kh=1, kw=1)), 44 | ("conv_1", Conv2dLayer(48, 64, kh=5, kw=5, padding=2)), 45 | ] 46 | ) 47 | ) 48 | self.tower_1 = torch.nn.Sequential( 49 | collections.OrderedDict( 50 | [ 51 | ("conv", Conv2dLayer(in_channels, 64, kh=1, kw=1)), 52 | ("conv_1", Conv2dLayer(64, 96, kh=3, kw=3, padding=1)), 53 | ("conv_2", Conv2dLayer(96, 96, kh=3, kw=3, padding=1)), 54 | ] 55 | ) 56 | ) 57 | self.tower_2 = torch.nn.Sequential( 58 | collections.OrderedDict( 59 | [ 60 | ( 61 | "pool", 62 | torch.nn.AvgPool2d( 63 | kernel_size=3, stride=1, padding=1, count_include_pad=False 64 | ), 65 | ), 66 | ("conv", Conv2dLayer(in_channels, tmp_channels, kh=1, kw=1)), 67 | ] 68 | ) 69 | ) 70 | 71 | def forward(self, x): 72 | return torch.cat( 73 | [ 74 | self.conv(x).contiguous(), 75 | self.tower(x).contiguous(), 76 | self.tower_1(x).contiguous(), 77 | self.tower_2(x).contiguous(), 78 | ], 79 | dim=1, 80 | ) 81 | 82 | 83 | # ---------------------------------------------------------------------------- 84 | 85 | 86 | class InceptionB(torch.nn.Module): 87 | def __init__(self, in_channels): 88 | super().__init__() 89 | self.conv = Conv2dLayer(in_channels, 384, kh=3, kw=3, stride=2) 90 | self.tower = torch.nn.Sequential( 91 | collections.OrderedDict( 92 | [ 93 | ("conv", Conv2dLayer(in_channels, 64, kh=1, kw=1)), 94 | ("conv_1", Conv2dLayer(64, 96, kh=3, kw=3, padding=1)), 95 | ("conv_2", Conv2dLayer(96, 96, kh=3, kw=3, stride=2)), 96 | ] 97 | ) 98 | ) 99 | self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=2) 100 | 101 | def forward(self, x): 102 | return torch.cat( 103 | [ 104 | self.conv(x).contiguous(), 105 | self.tower(x).contiguous(), 106 | self.pool(x).contiguous(), 107 | ], 108 | dim=1, 109 | ) 110 | 111 | 112 | # ---------------------------------------------------------------------------- 113 | 114 | 115 | class InceptionC(torch.nn.Module): 116 | def __init__(self, in_channels, tmp_channels): 117 | super().__init__() 118 | self.conv = Conv2dLayer(in_channels, 192, kh=1, kw=1) 119 | self.tower = torch.nn.Sequential( 120 | collections.OrderedDict( 121 | [ 122 | ("conv", Conv2dLayer(in_channels, tmp_channels, kh=1, kw=1)), 123 | ( 124 | "conv_1", 125 | Conv2dLayer( 126 | tmp_channels, tmp_channels, kh=1, kw=7, padding=[0, 3] 127 | ), 128 | ), 129 | ( 130 | "conv_2", 131 | Conv2dLayer(tmp_channels, 192, kh=7, kw=1, padding=[3, 0]), 132 | ), 133 | ] 134 | ) 135 | ) 136 | self.tower_1 = torch.nn.Sequential( 137 | collections.OrderedDict( 138 | [ 139 | ("conv", Conv2dLayer(in_channels, tmp_channels, kh=1, kw=1)), 140 | ( 141 | "conv_1", 142 | Conv2dLayer( 143 | tmp_channels, tmp_channels, kh=7, kw=1, padding=[3, 0] 144 | ), 145 | ), 146 | ( 147 | "conv_2", 148 | Conv2dLayer( 149 | tmp_channels, tmp_channels, kh=1, kw=7, padding=[0, 3] 150 | ), 151 | ), 152 | ( 153 | "conv_3", 154 | Conv2dLayer( 155 | tmp_channels, tmp_channels, kh=7, kw=1, padding=[3, 0] 156 | ), 157 | ), 158 | ( 159 | "conv_4", 160 | Conv2dLayer(tmp_channels, 192, kh=1, kw=7, padding=[0, 3]), 161 | ), 162 | ] 163 | ) 164 | ) 165 | self.tower_2 = torch.nn.Sequential( 166 | collections.OrderedDict( 167 | [ 168 | ( 169 | "pool", 170 | torch.nn.AvgPool2d( 171 | kernel_size=3, stride=1, padding=1, count_include_pad=False 172 | ), 173 | ), 174 | ("conv", Conv2dLayer(in_channels, 192, kh=1, kw=1)), 175 | ] 176 | ) 177 | ) 178 | 179 | def forward(self, x): 180 | return torch.cat( 181 | [ 182 | self.conv(x).contiguous(), 183 | self.tower(x).contiguous(), 184 | self.tower_1(x).contiguous(), 185 | self.tower_2(x).contiguous(), 186 | ], 187 | dim=1, 188 | ) 189 | 190 | 191 | # ---------------------------------------------------------------------------- 192 | 193 | 194 | class InceptionD(torch.nn.Module): 195 | def __init__(self, in_channels): 196 | super().__init__() 197 | self.tower = torch.nn.Sequential( 198 | collections.OrderedDict( 199 | [ 200 | ("conv", Conv2dLayer(in_channels, 192, kh=1, kw=1)), 201 | ("conv_1", Conv2dLayer(192, 320, kh=3, kw=3, stride=2)), 202 | ] 203 | ) 204 | ) 205 | self.tower_1 = torch.nn.Sequential( 206 | collections.OrderedDict( 207 | [ 208 | ("conv", Conv2dLayer(in_channels, 192, kh=1, kw=1)), 209 | ("conv_1", Conv2dLayer(192, 192, kh=1, kw=7, padding=[0, 3])), 210 | ("conv_2", Conv2dLayer(192, 192, kh=7, kw=1, padding=[3, 0])), 211 | ("conv_3", Conv2dLayer(192, 192, kh=3, kw=3, stride=2)), 212 | ] 213 | ) 214 | ) 215 | self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=2) 216 | 217 | def forward(self, x): 218 | return torch.cat( 219 | [ 220 | self.tower(x).contiguous(), 221 | self.tower_1(x).contiguous(), 222 | self.pool(x).contiguous(), 223 | ], 224 | dim=1, 225 | ) 226 | 227 | 228 | # ---------------------------------------------------------------------------- 229 | 230 | 231 | class InceptionE(torch.nn.Module): 232 | def __init__(self, in_channels, use_avg_pool): 233 | super().__init__() 234 | self.conv = Conv2dLayer(in_channels, 320, kh=1, kw=1) 235 | self.tower_conv = Conv2dLayer(in_channels, 384, kh=1, kw=1) 236 | self.tower_mixed_conv = Conv2dLayer(384, 384, kh=1, kw=3, padding=[0, 1]) 237 | self.tower_mixed_conv_1 = Conv2dLayer(384, 384, kh=3, kw=1, padding=[1, 0]) 238 | self.tower_1_conv = Conv2dLayer(in_channels, 448, kh=1, kw=1) 239 | self.tower_1_conv_1 = Conv2dLayer(448, 384, kh=3, kw=3, padding=1) 240 | self.tower_1_mixed_conv = Conv2dLayer(384, 384, kh=1, kw=3, padding=[0, 1]) 241 | self.tower_1_mixed_conv_1 = Conv2dLayer(384, 384, kh=3, kw=1, padding=[1, 0]) 242 | if use_avg_pool: 243 | self.tower_2_pool = torch.nn.AvgPool2d( 244 | kernel_size=3, stride=1, padding=1, count_include_pad=False 245 | ) 246 | else: 247 | self.tower_2_pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1) 248 | self.tower_2_conv = Conv2dLayer(in_channels, 192, kh=1, kw=1) 249 | 250 | def forward(self, x): 251 | a = self.tower_conv(x) 252 | b = self.tower_1_conv_1(self.tower_1_conv(x)) 253 | return torch.cat( 254 | [ 255 | self.conv(x).contiguous(), 256 | self.tower_mixed_conv(a).contiguous(), 257 | self.tower_mixed_conv_1(a).contiguous(), 258 | self.tower_1_mixed_conv(b).contiguous(), 259 | self.tower_1_mixed_conv_1(b).contiguous(), 260 | self.tower_2_conv(self.tower_2_pool(x)).contiguous(), 261 | ], 262 | dim=1, 263 | ) 264 | 265 | 266 | # ---------------------------------------------------------------------------- 267 | 268 | 269 | class InceptionV3(torch.nn.Module): 270 | def __init__(self): 271 | super().__init__() 272 | self.layers = torch.nn.Sequential( 273 | collections.OrderedDict( 274 | [ 275 | ("conv", Conv2dLayer(3, 32, kh=3, kw=3, stride=2)), 276 | ("conv_1", Conv2dLayer(32, 32, kh=3, kw=3)), 277 | ("conv_2", Conv2dLayer(32, 64, kh=3, kw=3, padding=1)), 278 | ("pool0", torch.nn.MaxPool2d(kernel_size=3, stride=2)), 279 | ("conv_3", Conv2dLayer(64, 80, kh=1, kw=1)), 280 | ("conv_4", Conv2dLayer(80, 192, kh=3, kw=3)), 281 | ("pool1", torch.nn.MaxPool2d(kernel_size=3, stride=2)), 282 | ("mixed", InceptionA(192, tmp_channels=32)), 283 | ("mixed_1", InceptionA(256, tmp_channels=64)), 284 | ("mixed_2", InceptionA(288, tmp_channels=64)), 285 | ("mixed_3", InceptionB(288)), 286 | ("mixed_4", InceptionC(768, tmp_channels=128)), 287 | ("mixed_5", InceptionC(768, tmp_channels=160)), 288 | ("mixed_6", InceptionC(768, tmp_channels=160)), 289 | ("mixed_7", InceptionC(768, tmp_channels=192)), 290 | ("mixed_8", InceptionD(768)), 291 | ("mixed_9", InceptionE(1280, use_avg_pool=True)), 292 | ("mixed_10", InceptionE(2048, use_avg_pool=False)), 293 | ("pool2", torch.nn.AvgPool2d(kernel_size=8)), 294 | ] 295 | ) 296 | ) 297 | self.output = torch.nn.Linear(2048, 1008) 298 | 299 | def forward( 300 | self, 301 | img, 302 | return_features: bool = True, 303 | use_fp16: bool = False, 304 | no_output_bias: bool = False, 305 | ): 306 | batch_size, channels, height, width = img.shape # [NCHW] 307 | assert channels == 3 308 | 309 | # Cast to float. 310 | x = img.to(torch.float16 if use_fp16 else torch.float32) 311 | 312 | # Emulate tf.image.resize_bilinear(x, [299, 299]), including the funky alignment. 313 | new_width, new_height = 299, 299 314 | theta = torch.eye(2, 3, device=x.device) 315 | theta[0, 2] += theta[0, 0] / width - theta[0, 0] / new_width 316 | theta[1, 2] += theta[1, 1] / height - theta[1, 1] / new_height 317 | theta = theta.to(x.dtype).unsqueeze(0).repeat([batch_size, 1, 1]) 318 | grid = torch.nn.functional.affine_grid( 319 | theta, [batch_size, channels, new_height, new_width], align_corners=False 320 | ) 321 | x = torch.nn.functional.grid_sample( 322 | x, grid, mode="bilinear", padding_mode="border", align_corners=False 323 | ) 324 | 325 | # Scale dynamic range from [0,255] to [-1,1[. 326 | x -= 128 327 | x /= 128 328 | 329 | # Main layers. 330 | intermediate = self.layers[:-6](x) 331 | spatial_features = ( 332 | self.layers[-6] 333 | .conv(intermediate)[:, :7] 334 | .permute(0, 2, 3, 1) 335 | .reshape(-1, 2023) 336 | ) 337 | features = self.layers[-6:](intermediate).reshape(-1, 2048).to(torch.float32) 338 | if return_features: 339 | return features, spatial_features 340 | 341 | # Output layer. 342 | return self.acts_to_probs(features, no_output_bias=no_output_bias) 343 | 344 | def acts_to_probs(self, features, no_output_bias: bool = False): 345 | if no_output_bias: 346 | logits = torch.nn.functional.linear(features, self.output.weight) 347 | else: 348 | logits = self.output(features) 349 | probs = torch.nn.functional.softmax(logits, dim=1) 350 | return probs 351 | 352 | def create_softmax_model(self): 353 | return SoftmaxModel(self.output.weight) 354 | 355 | 356 | class SoftmaxModel(torch.nn.Module): 357 | def __init__(self, weight: torch.Tensor): 358 | super().__init__() 359 | self.weight = torch.nn.Parameter(weight.detach().clone()) 360 | 361 | def forward(self, x): 362 | logits = torch.nn.functional.linear(x, self.weight) 363 | probs = torch.nn.functional.softmax(logits, dim=1) 364 | return probs 365 | -------------------------------------------------------------------------------- /iCT/evaluations/inception_v3.py: -------------------------------------------------------------------------------- 1 | # Ported from the model here: 2 | # https://github.com/NVlabs/stylegan3/blob/407db86e6fe432540a22515310188288687858fa/metrics/frechet_inception_distance.py#L22 3 | # 4 | # I have verified that the spatial features and output features are correct 5 | # within a mean absolute error of ~3e-5. 6 | 7 | import collections 8 | 9 | import torch 10 | 11 | 12 | class Conv2dLayer(torch.nn.Module): 13 | def __init__(self, in_channels, out_channels, kh, kw, stride=1, padding=0): 14 | super().__init__() 15 | self.stride = stride 16 | self.padding = padding 17 | self.weight = torch.nn.Parameter(torch.zeros(out_channels, in_channels, kh, kw)) 18 | self.beta = torch.nn.Parameter(torch.zeros(out_channels)) 19 | self.mean = torch.nn.Parameter(torch.zeros(out_channels)) 20 | self.var = torch.nn.Parameter(torch.zeros(out_channels)) 21 | 22 | def forward(self, x): 23 | x = torch.nn.functional.conv2d( 24 | x, self.weight.to(x.dtype), stride=self.stride, padding=self.padding 25 | ) 26 | x = torch.nn.functional.batch_norm( 27 | x, running_mean=self.mean, running_var=self.var, bias=self.beta, eps=1e-3 28 | ) 29 | x = torch.nn.functional.relu(x) 30 | return x 31 | 32 | 33 | # ---------------------------------------------------------------------------- 34 | 35 | 36 | class InceptionA(torch.nn.Module): 37 | def __init__(self, in_channels, tmp_channels): 38 | super().__init__() 39 | self.conv = Conv2dLayer(in_channels, 64, kh=1, kw=1) 40 | self.tower = torch.nn.Sequential( 41 | collections.OrderedDict( 42 | [ 43 | ("conv", Conv2dLayer(in_channels, 48, kh=1, kw=1)), 44 | ("conv_1", Conv2dLayer(48, 64, kh=5, kw=5, padding=2)), 45 | ] 46 | ) 47 | ) 48 | self.tower_1 = torch.nn.Sequential( 49 | collections.OrderedDict( 50 | [ 51 | ("conv", Conv2dLayer(in_channels, 64, kh=1, kw=1)), 52 | ("conv_1", Conv2dLayer(64, 96, kh=3, kw=3, padding=1)), 53 | ("conv_2", Conv2dLayer(96, 96, kh=3, kw=3, padding=1)), 54 | ] 55 | ) 56 | ) 57 | self.tower_2 = torch.nn.Sequential( 58 | collections.OrderedDict( 59 | [ 60 | ( 61 | "pool", 62 | torch.nn.AvgPool2d( 63 | kernel_size=3, stride=1, padding=1, count_include_pad=False 64 | ), 65 | ), 66 | ("conv", Conv2dLayer(in_channels, tmp_channels, kh=1, kw=1)), 67 | ] 68 | ) 69 | ) 70 | 71 | def forward(self, x): 72 | return torch.cat( 73 | [ 74 | self.conv(x).contiguous(), 75 | self.tower(x).contiguous(), 76 | self.tower_1(x).contiguous(), 77 | self.tower_2(x).contiguous(), 78 | ], 79 | dim=1, 80 | ) 81 | 82 | 83 | # ---------------------------------------------------------------------------- 84 | 85 | 86 | class InceptionB(torch.nn.Module): 87 | def __init__(self, in_channels): 88 | super().__init__() 89 | self.conv = Conv2dLayer(in_channels, 384, kh=3, kw=3, stride=2) 90 | self.tower = torch.nn.Sequential( 91 | collections.OrderedDict( 92 | [ 93 | ("conv", Conv2dLayer(in_channels, 64, kh=1, kw=1)), 94 | ("conv_1", Conv2dLayer(64, 96, kh=3, kw=3, padding=1)), 95 | ("conv_2", Conv2dLayer(96, 96, kh=3, kw=3, stride=2)), 96 | ] 97 | ) 98 | ) 99 | self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=2) 100 | 101 | def forward(self, x): 102 | return torch.cat( 103 | [ 104 | self.conv(x).contiguous(), 105 | self.tower(x).contiguous(), 106 | self.pool(x).contiguous(), 107 | ], 108 | dim=1, 109 | ) 110 | 111 | 112 | # ---------------------------------------------------------------------------- 113 | 114 | 115 | class InceptionC(torch.nn.Module): 116 | def __init__(self, in_channels, tmp_channels): 117 | super().__init__() 118 | self.conv = Conv2dLayer(in_channels, 192, kh=1, kw=1) 119 | self.tower = torch.nn.Sequential( 120 | collections.OrderedDict( 121 | [ 122 | ("conv", Conv2dLayer(in_channels, tmp_channels, kh=1, kw=1)), 123 | ( 124 | "conv_1", 125 | Conv2dLayer( 126 | tmp_channels, tmp_channels, kh=1, kw=7, padding=[0, 3] 127 | ), 128 | ), 129 | ( 130 | "conv_2", 131 | Conv2dLayer(tmp_channels, 192, kh=7, kw=1, padding=[3, 0]), 132 | ), 133 | ] 134 | ) 135 | ) 136 | self.tower_1 = torch.nn.Sequential( 137 | collections.OrderedDict( 138 | [ 139 | ("conv", Conv2dLayer(in_channels, tmp_channels, kh=1, kw=1)), 140 | ( 141 | "conv_1", 142 | Conv2dLayer( 143 | tmp_channels, tmp_channels, kh=7, kw=1, padding=[3, 0] 144 | ), 145 | ), 146 | ( 147 | "conv_2", 148 | Conv2dLayer( 149 | tmp_channels, tmp_channels, kh=1, kw=7, padding=[0, 3] 150 | ), 151 | ), 152 | ( 153 | "conv_3", 154 | Conv2dLayer( 155 | tmp_channels, tmp_channels, kh=7, kw=1, padding=[3, 0] 156 | ), 157 | ), 158 | ( 159 | "conv_4", 160 | Conv2dLayer(tmp_channels, 192, kh=1, kw=7, padding=[0, 3]), 161 | ), 162 | ] 163 | ) 164 | ) 165 | self.tower_2 = torch.nn.Sequential( 166 | collections.OrderedDict( 167 | [ 168 | ( 169 | "pool", 170 | torch.nn.AvgPool2d( 171 | kernel_size=3, stride=1, padding=1, count_include_pad=False 172 | ), 173 | ), 174 | ("conv", Conv2dLayer(in_channels, 192, kh=1, kw=1)), 175 | ] 176 | ) 177 | ) 178 | 179 | def forward(self, x): 180 | return torch.cat( 181 | [ 182 | self.conv(x).contiguous(), 183 | self.tower(x).contiguous(), 184 | self.tower_1(x).contiguous(), 185 | self.tower_2(x).contiguous(), 186 | ], 187 | dim=1, 188 | ) 189 | 190 | 191 | # ---------------------------------------------------------------------------- 192 | 193 | 194 | class InceptionD(torch.nn.Module): 195 | def __init__(self, in_channels): 196 | super().__init__() 197 | self.tower = torch.nn.Sequential( 198 | collections.OrderedDict( 199 | [ 200 | ("conv", Conv2dLayer(in_channels, 192, kh=1, kw=1)), 201 | ("conv_1", Conv2dLayer(192, 320, kh=3, kw=3, stride=2)), 202 | ] 203 | ) 204 | ) 205 | self.tower_1 = torch.nn.Sequential( 206 | collections.OrderedDict( 207 | [ 208 | ("conv", Conv2dLayer(in_channels, 192, kh=1, kw=1)), 209 | ("conv_1", Conv2dLayer(192, 192, kh=1, kw=7, padding=[0, 3])), 210 | ("conv_2", Conv2dLayer(192, 192, kh=7, kw=1, padding=[3, 0])), 211 | ("conv_3", Conv2dLayer(192, 192, kh=3, kw=3, stride=2)), 212 | ] 213 | ) 214 | ) 215 | self.pool = torch.nn.MaxPool2d(kernel_size=3, stride=2) 216 | 217 | def forward(self, x): 218 | return torch.cat( 219 | [ 220 | self.tower(x).contiguous(), 221 | self.tower_1(x).contiguous(), 222 | self.pool(x).contiguous(), 223 | ], 224 | dim=1, 225 | ) 226 | 227 | 228 | # ---------------------------------------------------------------------------- 229 | 230 | 231 | class InceptionE(torch.nn.Module): 232 | def __init__(self, in_channels, use_avg_pool): 233 | super().__init__() 234 | self.conv = Conv2dLayer(in_channels, 320, kh=1, kw=1) 235 | self.tower_conv = Conv2dLayer(in_channels, 384, kh=1, kw=1) 236 | self.tower_mixed_conv = Conv2dLayer(384, 384, kh=1, kw=3, padding=[0, 1]) 237 | self.tower_mixed_conv_1 = Conv2dLayer(384, 384, kh=3, kw=1, padding=[1, 0]) 238 | self.tower_1_conv = Conv2dLayer(in_channels, 448, kh=1, kw=1) 239 | self.tower_1_conv_1 = Conv2dLayer(448, 384, kh=3, kw=3, padding=1) 240 | self.tower_1_mixed_conv = Conv2dLayer(384, 384, kh=1, kw=3, padding=[0, 1]) 241 | self.tower_1_mixed_conv_1 = Conv2dLayer(384, 384, kh=3, kw=1, padding=[1, 0]) 242 | if use_avg_pool: 243 | self.tower_2_pool = torch.nn.AvgPool2d( 244 | kernel_size=3, stride=1, padding=1, count_include_pad=False 245 | ) 246 | else: 247 | self.tower_2_pool = torch.nn.MaxPool2d(kernel_size=3, stride=1, padding=1) 248 | self.tower_2_conv = Conv2dLayer(in_channels, 192, kh=1, kw=1) 249 | 250 | def forward(self, x): 251 | a = self.tower_conv(x) 252 | b = self.tower_1_conv_1(self.tower_1_conv(x)) 253 | return torch.cat( 254 | [ 255 | self.conv(x).contiguous(), 256 | self.tower_mixed_conv(a).contiguous(), 257 | self.tower_mixed_conv_1(a).contiguous(), 258 | self.tower_1_mixed_conv(b).contiguous(), 259 | self.tower_1_mixed_conv_1(b).contiguous(), 260 | self.tower_2_conv(self.tower_2_pool(x)).contiguous(), 261 | ], 262 | dim=1, 263 | ) 264 | 265 | 266 | # ---------------------------------------------------------------------------- 267 | 268 | 269 | class InceptionV3(torch.nn.Module): 270 | def __init__(self): 271 | super().__init__() 272 | self.layers = torch.nn.Sequential( 273 | collections.OrderedDict( 274 | [ 275 | ("conv", Conv2dLayer(3, 32, kh=3, kw=3, stride=2)), 276 | ("conv_1", Conv2dLayer(32, 32, kh=3, kw=3)), 277 | ("conv_2", Conv2dLayer(32, 64, kh=3, kw=3, padding=1)), 278 | ("pool0", torch.nn.MaxPool2d(kernel_size=3, stride=2)), 279 | ("conv_3", Conv2dLayer(64, 80, kh=1, kw=1)), 280 | ("conv_4", Conv2dLayer(80, 192, kh=3, kw=3)), 281 | ("pool1", torch.nn.MaxPool2d(kernel_size=3, stride=2)), 282 | ("mixed", InceptionA(192, tmp_channels=32)), 283 | ("mixed_1", InceptionA(256, tmp_channels=64)), 284 | ("mixed_2", InceptionA(288, tmp_channels=64)), 285 | ("mixed_3", InceptionB(288)), 286 | ("mixed_4", InceptionC(768, tmp_channels=128)), 287 | ("mixed_5", InceptionC(768, tmp_channels=160)), 288 | ("mixed_6", InceptionC(768, tmp_channels=160)), 289 | ("mixed_7", InceptionC(768, tmp_channels=192)), 290 | ("mixed_8", InceptionD(768)), 291 | ("mixed_9", InceptionE(1280, use_avg_pool=True)), 292 | ("mixed_10", InceptionE(2048, use_avg_pool=False)), 293 | ("pool2", torch.nn.AvgPool2d(kernel_size=8)), 294 | ] 295 | ) 296 | ) 297 | self.output = torch.nn.Linear(2048, 1008) 298 | 299 | def forward( 300 | self, 301 | img, 302 | return_features: bool = True, 303 | use_fp16: bool = False, 304 | no_output_bias: bool = False, 305 | ): 306 | batch_size, channels, height, width = img.shape # [NCHW] 307 | assert channels == 3 308 | 309 | # Cast to float. 310 | x = img.to(torch.float16 if use_fp16 else torch.float32) 311 | 312 | # Emulate tf.image.resize_bilinear(x, [299, 299]), including the funky alignment. 313 | new_width, new_height = 299, 299 314 | theta = torch.eye(2, 3, device=x.device) 315 | theta[0, 2] += theta[0, 0] / width - theta[0, 0] / new_width 316 | theta[1, 2] += theta[1, 1] / height - theta[1, 1] / new_height 317 | theta = theta.to(x.dtype).unsqueeze(0).repeat([batch_size, 1, 1]) 318 | grid = torch.nn.functional.affine_grid( 319 | theta, [batch_size, channels, new_height, new_width], align_corners=False 320 | ) 321 | x = torch.nn.functional.grid_sample( 322 | x, grid, mode="bilinear", padding_mode="border", align_corners=False 323 | ) 324 | 325 | # Scale dynamic range from [0,255] to [-1,1[. 326 | x -= 128 327 | x /= 128 328 | 329 | # Main layers. 330 | intermediate = self.layers[:-6](x) 331 | spatial_features = ( 332 | self.layers[-6] 333 | .conv(intermediate)[:, :7] 334 | .permute(0, 2, 3, 1) 335 | .reshape(-1, 2023) 336 | ) 337 | features = self.layers[-6:](intermediate).reshape(-1, 2048).to(torch.float32) 338 | if return_features: 339 | return features, spatial_features 340 | 341 | # Output layer. 342 | return self.acts_to_probs(features, no_output_bias=no_output_bias) 343 | 344 | def acts_to_probs(self, features, no_output_bias: bool = False): 345 | if no_output_bias: 346 | logits = torch.nn.functional.linear(features, self.output.weight) 347 | else: 348 | logits = self.output(features) 349 | probs = torch.nn.functional.softmax(logits, dim=1) 350 | return probs 351 | 352 | def create_softmax_model(self): 353 | return SoftmaxModel(self.output.weight) 354 | 355 | 356 | class SoftmaxModel(torch.nn.Module): 357 | def __init__(self, weight: torch.Tensor): 358 | super().__init__() 359 | self.weight = torch.nn.Parameter(weight.detach().clone()) 360 | 361 | def forward(self, x): 362 | logits = torch.nn.functional.linear(x, self.weight) 363 | probs = torch.nn.functional.softmax(logits, dim=1) 364 | return probs 365 | -------------------------------------------------------------------------------- /BCM/cm/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 | --------------------------------------------------------------------------------